<commit_before>package slack

import (
	"encoding/json"
	"errors"
	"net/url"
	"strings"
)

const (
	DEFAULT_MESSAGE_USERNAME     = ""
	DEFAULT_MESSAGE_ASUSER       = false
	DEFAULT_MESSAGE_PARSE        = ""
	DEFAULT_MESSAGE_LINK_NAMES   = 0
	DEFAULT_MESSAGE_UNFURL_LINKS = true
	DEFAULT_MESSAGE_UNFURL_MEDIA = false
	DEFAULT_MESSAGE_ICON_URL     = ""
	DEFAULT_MESSAGE_ICON_EMOJI   = ""
	DEFAULT_MESSAGE_MARKDOWN     = true
	DEFAULT_MESSAGE_ESCAPE_TEXT  = true
)

type chatResponseFull struct {
	ChannelId string `json:"channel"`
	Timestamp string `json:"ts"`
	Text      string `json:"text"`
	SlackResponse
}

// AttachmentField contains information for an attachment field
// An Attachment can contain multiple of these
type AttachmentField struct {
	Title string `json:"title"`
	Value string `json:"value"`
	Short bool   `json:"short"`
}

// Attachment contains all the information for an attachment
type Attachment struct {
	Fallback string `json:"fallback"`

	Color string `json:"color,omitempty"`

	Pretext string `json:"pretext,omitempty"`

	AuthorName string `json:"author_name,omitempty"`
	AuthorLink string `json:"author_link,omitempty"`
	AuthorIcon string `json:"author_icon,omitempty"`

	Title     string `json:"title,omitempty"`
	TitleLink string `json:"title_link,omitempty"`

	Text string `json:"text"`

	ImageURL string `json:"image_url,omitempty"`

	Fields []AttachmentField `json:"fields,omitempty"`

	MarkdownIn []string `json:"mrkdwn_in,omitempty"`
}

// PostMessageParameters contains all the parameters necessary (including the optional ones) for a PostMessage() request
type PostMessageParameters struct {
	Text        string
	Username    string
	AsUser      bool
	Parse       string
	LinkNames   int
	Attachments []Attachment
	UnfurlLinks bool
	UnfurlMedia bool
	IconURL     string
	IconEmoji   string
	Markdown    bool `json:"mrkdwn,omitempty"`
	EscapeText  bool
}

// NewPostMessageParameters provides an instance of PostMessageParameters with all the sane default values set
func NewPostMessageParameters() PostMessageParameters {
	return PostMessageParameters{
		Username:    DEFAULT_MESSAGE_USERNAME,
		AsUser:      DEFAULT_MESSAGE_ASUSER,
		Parse:       DEFAULT_MESSAGE_PARSE,
		LinkNames:   DEFAULT_MESSAGE_LINK_NAMES,
		Attachments: nil,
		UnfurlLinks: DEFAULT_MESSAGE_UNFURL_LINKS,
		UnfurlMedia: DEFAULT_MESSAGE_UNFURL_MEDIA,
		IconURL:     DEFAULT_MESSAGE_ICON_URL,
		IconEmoji:   DEFAULT_MESSAGE_ICON_EMOJI,
		Markdown:    DEFAULT_MESSAGE_MARKDOWN,
		EscapeText:  DEFAULT_MESSAGE_ESCAPE_TEXT,
	}
}

func chatRequest(path string, values url.Values, debug bool) (*chatResponseFull, error) {
	response := &chatResponseFull{}
	err := parseResponse(path, values, response, debug)
	if err != nil {
		return nil, err
	}
	if !response.Ok {
		return nil, errors.New(response.Error)
	}
	return response, nil
}

// DeleteMessage deletes a message in a channel
func (api *Slack) DeleteMessage(channelId, messageTimestamp string) (string, string, error) {
	values := url.Values{
		"token":   {api.config.token},
		"channel": {channelId},
		"ts":      {messageTimestamp},
	}
	response, err := chatRequest("chat.delete", values, api.debug)
	if err != nil {
		return "", "", err
	}
	return response.ChannelId, response.Timestamp, nil
}

func escapeMessage(message string) string {
	/*
		& replaced with &
		< replaced with <
		> replaced with >
	*/
	replacer := strings.NewReplacer("&", "&", "<", "<", ">", ">")
	return replacer.Replace(message)
}

// PostMessage sends a message to a channel
// Message is escaped by default according to https://api.slack.com/docs/formatting
func (api *Slack) PostMessage(channelId string, text string, params PostMessageParameters) (channel string, timestamp string, err error) {
	if params.EscapeText {
		text = escapeMessage(text)
	}
	values := url.Values{
		"token":   {api.config.token},
		"channel": {channelId},
		"text":    {text},
	}
	if params.Username != DEFAULT_MESSAGE_USERNAME {
		values.Set("username", string(params.Username))
	}
	if params.AsUser != DEFAULT_MESSAGE_ASUSER {
		values.Set("as_user", "true")
	}
	if params.Parse != DEFAULT_MESSAGE_PARSE {
		values.Set("parse", string(params.Parse))
	}
	if params.LinkNames != DEFAULT_MESSAGE_LINK_NAMES {
		values.Set("link_names", "1")
	}
	if params.Attachments != nil {
		attachments, err := json.Marshal(params.Attachments)
		if err != nil {
			return "", "", err
		}
		values.Set("attachments", string(attachments))
	}
	if params.UnfurlLinks == DEFAULT_MESSAGE_UNFURL_LINKS {
		values.Set("unfurl_links", "false")
	}
	if params.UnfurlMedia != DEFAULT_MESSAGE_UNFURL_MEDIA {
		values.Set("unfurl_media", "true")
	}
	if params.IconURL != DEFAULT_MESSAGE_ICON_URL {
		values.Set("icon_url", params.IconURL)
	}
	if params.IconEmoji != DEFAULT_MESSAGE_ICON_EMOJI {
		values.Set("icon_emoji", params.IconEmoji)
	}
	if params.Markdown != DEFAULT_MESSAGE_MARKDOWN {
		values.Set("mrkdwn", "false")
	}

	response, err := chatRequest("chat.postMessage", values, api.debug)
	if err != nil {
		return "", "", err
	}
	return response.ChannelId, response.Timestamp, nil
}

// UpdateMessage updates a message in a channel
func (api *Slack) UpdateMessage(channelId, timestamp, text string) (string, string, string, error) {
	values := url.Values{
		"token":   {api.config.token},
		"channel": {channelId},
		"text":    {escapeMessage(text)},
		"ts":      {timestamp},
	}
	response, err := chatRequest("chat.update", values, api.debug)
	if err != nil {
		return "", "", "", err
	}
	return response.ChannelId, response.Timestamp, response.Text, nil
}
<commit_msg>fix unfurl_links does not work.<commit_after>package slack

import (
	"encoding/json"
	"errors"
	"net/url"
	"strings"
)

const (
	DEFAULT_MESSAGE_USERNAME     = ""
	DEFAULT_MESSAGE_ASUSER       = false
	DEFAULT_MESSAGE_PARSE        = ""
	DEFAULT_MESSAGE_LINK_NAMES   = 0
	DEFAULT_MESSAGE_UNFURL_LINKS = false
	DEFAULT_MESSAGE_UNFURL_MEDIA = false
	DEFAULT_MESSAGE_ICON_URL     = ""
	DEFAULT_MESSAGE_ICON_EMOJI   = ""
	DEFAULT_MESSAGE_MARKDOWN     = true
	DEFAULT_MESSAGE_ESCAPE_TEXT  = true
)

type chatResponseFull struct {
	ChannelId string `json:"channel"`
	Timestamp string `json:"ts"`
	Text      string `json:"text"`
	SlackResponse
}

// AttachmentField contains information for an attachment field
// An Attachment can contain multiple of these
type AttachmentField struct {
	Title string `json:"title"`
	Value string `json:"value"`
	Short bool   `json:"short"`
}

// Attachment contains all the information for an attachment
type Attachment struct {
	Fallback string `json:"fallback"`

	Color string `json:"color,omitempty"`

	Pretext string `json:"pretext,omitempty"`

	AuthorName string `json:"author_name,omitempty"`
	AuthorLink string `json:"author_link,omitempty"`
	AuthorIcon string `json:"author_icon,omitempty"`

	Title     string `json:"title,omitempty"`
	TitleLink string `json:"title_link,omitempty"`

	Text string `json:"text"`

	ImageURL string `json:"image_url,omitempty"`

	Fields []AttachmentField `json:"fields,omitempty"`

	MarkdownIn []string `json:"mrkdwn_in,omitempty"`
}

// PostMessageParameters contains all the parameters necessary (including the optional ones) for a PostMessage() request
type PostMessageParameters struct {
	Text        string
	Username    string
	AsUser      bool
	Parse       string
	LinkNames   int
	Attachments []Attachment
	UnfurlLinks bool
	UnfurlMedia bool
	IconURL     string
	IconEmoji   string
	Markdown    bool `json:"mrkdwn,omitempty"`
	EscapeText  bool
}

// NewPostMessageParameters provides an instance of PostMessageParameters with all the sane default values set
func NewPostMessageParameters() PostMessageParameters {
	return PostMessageParameters{
		Username:    DEFAULT_MESSAGE_USERNAME,
		AsUser:      DEFAULT_MESSAGE_ASUSER,
		Parse:       DEFAULT_MESSAGE_PARSE,
		LinkNames:   DEFAULT_MESSAGE_LINK_NAMES,
		Attachments: nil,
		UnfurlLinks: DEFAULT_MESSAGE_UNFURL_LINKS,
		UnfurlMedia: DEFAULT_MESSAGE_UNFURL_MEDIA,
		IconURL:     DEFAULT_MESSAGE_ICON_URL,
		IconEmoji:   DEFAULT_MESSAGE_ICON_EMOJI,
		Markdown:    DEFAULT_MESSAGE_MARKDOWN,
		EscapeText:  DEFAULT_MESSAGE_ESCAPE_TEXT,
	}
}

func chatRequest(path string, values url.Values, debug bool) (*chatResponseFull, error) {
	response := &chatResponseFull{}
	err := parseResponse(path, values, response, debug)
	if err != nil {
		return nil, err
	}
	if !response.Ok {
		return nil, errors.New(response.Error)
	}
	return response, nil
}

// DeleteMessage deletes a message in a channel
func (api *Slack) DeleteMessage(channelId, messageTimestamp string) (string, string, error) {
	values := url.Values{
		"token":   {api.config.token},
		"channel": {channelId},
		"ts":      {messageTimestamp},
	}
	response, err := chatRequest("chat.delete", values, api.debug)
	if err != nil {
		return "", "", err
	}
	return response.ChannelId, response.Timestamp, nil
}

func escapeMessage(message string) string {
	/*
		& replaced with &
		< replaced with <
		> replaced with >
	*/
	replacer := strings.NewReplacer("&", "&", "<", "<", ">", ">")
	return replacer.Replace(message)
}

// PostMessage sends a message to a channel
// Message is escaped by default according to https://api.slack.com/docs/formatting
func (api *Slack) PostMessage(channelId string, text string, params PostMessageParameters) (channel string, timestamp string, err error) {
	if params.EscapeText {
		text = escapeMessage(text)
	}
	values := url.Values{
		"token":   {api.config.token},
		"channel": {channelId},
		"text":    {text},
	}
	if params.Username != DEFAULT_MESSAGE_USERNAME {
		values.Set("username", string(params.Username))
	}
	if params.AsUser != DEFAULT_MESSAGE_ASUSER {
		values.Set("as_user", "true")
	}
	if params.Parse != DEFAULT_MESSAGE_PARSE {
		values.Set("parse", string(params.Parse))
	}
	if params.LinkNames != DEFAULT_MESSAGE_LINK_NAMES {
		values.Set("link_names", "1")
	}
	if params.Attachments != nil {
		attachments, err := json.Marshal(params.Attachments)
		if err != nil {
			return "", "", err
		}
		values.Set("attachments", string(attachments))
	}
	if params.UnfurlLinks != DEFAULT_MESSAGE_UNFURL_LINKS {
		values.Set("unfurl_links", "true")
	}
	if params.UnfurlMedia != DEFAULT_MESSAGE_UNFURL_MEDIA {
		values.Set("unfurl_media", "true")
	}
	if params.IconURL != DEFAULT_MESSAGE_ICON_URL {
		values.Set("icon_url", params.IconURL)
	}
	if params.IconEmoji != DEFAULT_MESSAGE_ICON_EMOJI {
		values.Set("icon_emoji", params.IconEmoji)
	}
	if params.Markdown != DEFAULT_MESSAGE_MARKDOWN {
		values.Set("mrkdwn", "false")
	}

	response, err := chatRequest("chat.postMessage", values, api.debug)
	if err != nil {
		return "", "", err
	}
	return response.ChannelId, response.Timestamp, nil
}

// UpdateMessage updates a message in a channel
func (api *Slack) UpdateMessage(channelId, timestamp, text string) (string, string, string, error) {
	values := url.Values{
		"token":   {api.config.token},
		"channel": {channelId},
		"text":    {escapeMessage(text)},
		"ts":      {timestamp},
	}
	response, err := chatRequest("chat.update", values, api.debug)
	if err != nil {
		return "", "", "", err
	}
	return response.ChannelId, response.Timestamp, response.Text, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2011 Jacob Amrany
 *
 * 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
 *
 * http://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 main

import (
	"bufio"
	"code.google.com/p/gorilla/mux"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

func Min(a ...int) int {
	min := int(^uint(0) >> 1) // largest int
	for _, i := range a {
		if i < min {
			min = i
		}
	}
	return min
}
func Max(a ...int) int {
	max := int(0)
	for _, i := range a {
		if i > max {
			max = i
		}
	}
	return max
}

var iIndex *InvertedIndex
var fIndex *ForwardIndex
var corpusPath string

func main() {
	flag.StringVar(&corpusPath, "Corpus_File_Path", "w1_fixed.txt", "The path to the corpus file.  A file with terms separated by \n")
	var port string
	flag.StringVar(&port, "port", "8080", "The port you want the web call to listen on.")

	iIndex = NewInvertedIndex()
	fIndex = NewForwardIndex()

	InitIndex(iIndex, fIndex)

	r := mux.NewRouter()
	r.HandleFunc("/cleo/{query}", Search)
	http.Handle("/", r)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

//Search handles the web requests and writes the output as
//json data.
func Search(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	query := vars["query"]

	searchResult := CleoSearch(iIndex, fIndex, query)
	sort.Sort(ByScore{searchResult})
	myJson, _ := json.Marshal(searchResult)
	fmt.Fprintf(w, string(myJson))
}

func InitIndex(iIndex *InvertedIndex, fIndex *ForwardIndex) {
	//Read corpus
	file, _ := os.Open(corpusPath)

	r := bufio.NewReader(file)
	docID := 1

	for {
		line, err := r.ReadString('\n')
		if err != nil {
			break
		}
		filter := computeBloomFilter(line)

		iIndex.AddDoc(docID, line, filter) //insert into inverted index
		fIndex.AddDoc(docID, line)         //Insert into forward index

		docID++
	}
}

type RankedResults []RankedResult
type ByScore struct{ RankedResults }

func (s RankedResults) Len() int      { return len(s) }
func (s RankedResults) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s ByScore) Less(i, j int) bool  { return s.RankedResults[i].Score > s.RankedResults[j].Score }

type RankedResult struct {
	Word  string
	Score float64
}

//This is the meat of the search.  It first checks the inverted index
//for matches, then filters the potentially numerous results using
//the bloom filter.  Finally, it ranks the word using a Levenshtein
//distance.
func CleoSearch(iIndex *InvertedIndex, fIndex *ForwardIndex, query string) []RankedResult {
	t0 := time.Now()
	rslt := make([]RankedResult, 0, 0)
	fmt.Println("Query:", query)

	candidates := iIndex.Search(query) //First get candidates from Inverted Index
	qBloom := computeBloomFilter(query)

	for _, i := range candidates {
		if TestBytesFromQuery(i.bloom, qBloom) == true { //Filter using Bloom Filter
			c := fIndex.itemAt(i.docId) //Get whole document from Forward Index
			score := Score(query, c)    //Score the Forward Index between 0-1
			ranked := RankedResult{c, score}
			rslt = append(rslt, ranked)
		}
	}
	t1 := time.Now()
	fmt.Printf("The call took %v to run.\n", t1.Sub(t0))
	return rslt
}

//Iterates through all of the 8 bytes (64 bits) and tests
//each bit that is set to 1 in the query's filter against
//the bit in the comparison's filter.  If the bit is not
// also 1, you do not have a match.
func TestBytesFromQuery(bf int, qBloom int) bool {
	for i := uint(0); i < 64; i++ {
		//a & (1 << idx) == b & (1 << idx)
		if (bf&(1<<i) != (1 << i)) && qBloom&(1<<i) == (1<<i) {
			return false
		}
	}
	return true
}

func Score(query, candidate string) float64 {
	lev := LevenshteinDistance(query, candidate)
	length := Max(len(candidate), len(query))
	return float64(length-lev) / float64(length+lev) //Jacard score
}

//Levenshtein distance is the number of inserts, deletions,
//and substitutions that differentiate one word from another.
//This algorithm is dynamic programming found at
//http://en.wikipedia.org/wiki/Levenshtein_distance
func LevenshteinDistance(s, t string) int {
	m := len(s)
	n := len(t)
	width := n - 1
	d := make([]int, m*n)
	//y * w + h for position in array
	for i := 1; i < m; i++ {
		d[i*width+0] = i
	}

	for j := 1; j < n; j++ {
		d[0*width+j] = j
	}

	for j := 1; j < n; j++ {
		for i := 1; i < m; i++ {
			if s[i] == t[j] {
				d[i*width+j] = d[(i-1)*width+(j-1)]
			} else {
				d[i*width+j] = Min(d[(i-1)*width+j]+1, //deletion
					d[i*width+(j-1)]+1,     //insertion
					d[(i-1)*width+(j-1)]+1) //substitution
			}
		}
	}
	return d[m*(width)+0]
}

func getPrefix(query string) string {
	qLen := Min(len(query), 4)
	q := query[0:qLen]
	return strings.ToLower(q)
}

type Document struct {
	docId int
	bloom int
}

//Used for the bloom filter
const (
	FNV_BASIS_64 = uint64(14695981039346656037)
	FNV_PRIME_64 = uint64((1 << 40) + 435)
	FNV_MASK_64  = uint64(^uint64(0) >> 1)
	NUM_BITS     = 64

	FNV_BASIS_32 = uint32(0x811c9dc5)
	FNV_PRIME_32 = uint32((1 << 24) + 403)
	FNV_MASK_32  = uint32(^uint32(0) >> 1)
)

//The bloom filter of a word is 8 bytes in length
//and has each character added separately
func computeBloomFilter(s string) int {
	cnt := len(s)

	if cnt <= 0 {
		return 0
	}

	var filter int
	hash := uint64(0)

	for i := 0; i < cnt; i++ {
		c := s[i]

		//first hash function
		hash ^= uint64(0xFF & c)
		hash *= FNV_PRIME_64

		//second hash function (reduces collisions for bloom)
		hash ^= uint64(0xFF & (c >> 16))
		hash *= FNV_PRIME_64

		//position of the bit mod the number of bits (8 bytes = 64 bits)
		bitpos := hash % NUM_BITS
		if bitpos < 0 {
			bitpos += NUM_BITS
		}
		filter = filter | (1 << bitpos)
	}

	return filter
}

//Inverted Index - Maps the query prefix to the matching documents
type InvertedIndex map[string][]Document

func NewInvertedIndex() *InvertedIndex {
	i := make(InvertedIndex)
	return &i
}

func (x *InvertedIndex) Size() int {
	return len(map[string][]Document(*x))
}

func (x *InvertedIndex) AddDoc(docId int, doc string, bloom int) {
	for _, word := range strings.Fields(doc) {
		word = getPrefix(word)

		ref, ok := (*x)[word]
		if !ok {
			ref = nil
		}

		(*x)[word] = append(ref, Document{docId: docId, bloom: bloom})
	}
}

func (x *InvertedIndex) Search(query string) []Document {
	q := getPrefix(query)

	ref, ok := (*x)[q]

	if ok {
		return ref
	}
	return nil
}

//Forward Index - Maps the document id to the document
type ForwardIndex map[int]string

func NewForwardIndex() *ForwardIndex {
	i := make(ForwardIndex)
	return &i
}
func (x *ForwardIndex) AddDoc(docId int, doc string) {
	for _, word := range strings.Fields(doc) {
		_, ok := (*x)[docId]
		if !ok {
			(*x)[docId] = word
		}
	}
}
func (x *ForwardIndex) itemAt(i int) string {
	return (*x)[i]
}
<commit_msg>made into library<commit_after>/*
 * Copyright (c) 2011 Jacob Amrany
 *
 * 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
 *
 * http://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 cleo

import (
	"bufio"
	"code.google.com/p/gorilla/mux"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

func Min(a ...int) int {
	min := int(^uint(0) >> 1) // largest int
	for _, i := range a {
		if i < min {
			min = i
		}
	}
	return min
}
func Max(a ...int) int {
	max := int(0)
	for _, i := range a {
		if i > max {
			max = i
		}
	}
	return max
}

type indexContainer struct {
	iIndex *InvertedIndex
	fIndex *ForwardIndex
}

var m *indexContainer

func InitAndRun(corpusPath, port string) {
	m = &indexContainer{}
	m.iIndex = NewInvertedIndex()
	m.fIndex = NewForwardIndex()

	InitIndex(m.iIndex, m.fIndex, corpusPath)

	r := mux.NewRouter()
	r.HandleFunc("/cleo/{query}", Search)
	http.Handle("/", r)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

//Search handles the web requests and writes the output as
//json data.
func Search(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	query := vars["query"]

	searchResult := CleoSearch(m.iIndex, m.fIndex, query)
	sort.Sort(ByScore{searchResult})
	myJson, _ := json.Marshal(searchResult)
	fmt.Fprintf(w, string(myJson))
}

func InitIndex(iIndex *InvertedIndex, fIndex *ForwardIndex, corpusPath string) {
	//Read corpus
	file, _ := os.Open(corpusPath)

	r := bufio.NewReader(file)
	docID := 1

	for {
		line, err := r.ReadString('\n')
		if err != nil {
			break
		}
		filter := computeBloomFilter(line)

		iIndex.AddDoc(docID, line, filter) //insert into inverted index
		fIndex.AddDoc(docID, line)         //Insert into forward index

		docID++
	}
}

type RankedResults []RankedResult
type ByScore struct{ RankedResults }

func (s RankedResults) Len() int      { return len(s) }
func (s RankedResults) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s ByScore) Less(i, j int) bool  { return s.RankedResults[i].Score > s.RankedResults[j].Score }

type RankedResult struct {
	Word  string
	Score float64
}

//This is the meat of the search.  It first checks the inverted index
//for matches, then filters the potentially numerous results using
//the bloom filter.  Finally, it ranks the word using a Levenshtein
//distance.
func CleoSearch(iIndex *InvertedIndex, fIndex *ForwardIndex, query string) []RankedResult {
	t0 := time.Now()
	rslt := make([]RankedResult, 0, 0)
	fmt.Println("Query:", query)

	candidates := iIndex.Search(query) //First get candidates from Inverted Index
	qBloom := computeBloomFilter(query)

	for _, i := range candidates {
		if TestBytesFromQuery(i.bloom, qBloom) == true { //Filter using Bloom Filter
			c := fIndex.itemAt(i.docId) //Get whole document from Forward Index
			score := Score(query, c)    //Score the Forward Index between 0-1
			ranked := RankedResult{c, score}
			rslt = append(rslt, ranked)
		}
	}
	t1 := time.Now()
	fmt.Printf("The call took %v to run.\n", t1.Sub(t0))
	return rslt
}

//Iterates through all of the 8 bytes (64 bits) and tests
//each bit that is set to 1 in the query's filter against
//the bit in the comparison's filter.  If the bit is not
// also 1, you do not have a match.
func TestBytesFromQuery(bf int, qBloom int) bool {
	for i := uint(0); i < 64; i++ {
		//a & (1 << idx) == b & (1 << idx)
		if (bf&(1<<i) != (1 << i)) && qBloom&(1<<i) == (1<<i) {
			return false
		}
	}
	return true
}

func Score(query, candidate string) float64 {
	lev := LevenshteinDistance(query, candidate)
	length := Max(len(candidate), len(query))
	return float64(length-lev) / float64(length+lev) //Jacard score
}

//Levenshtein distance is the number of inserts, deletions,
//and substitutions that differentiate one word from another.
//This algorithm is dynamic programming found at
//http://en.wikipedia.org/wiki/Levenshtein_distance
func LevenshteinDistance(s, t string) int {
	m := len(s)
	n := len(t)
	width := n - 1
	d := make([]int, m*n)
	//y * w + h for position in array
	for i := 1; i < m; i++ {
		d[i*width+0] = i
	}

	for j := 1; j < n; j++ {
		d[0*width+j] = j
	}

	for j := 1; j < n; j++ {
		for i := 1; i < m; i++ {
			if s[i] == t[j] {
				d[i*width+j] = d[(i-1)*width+(j-1)]
			} else {
				d[i*width+j] = Min(d[(i-1)*width+j]+1, //deletion
					d[i*width+(j-1)]+1,     //insertion
					d[(i-1)*width+(j-1)]+1) //substitution
			}
		}
	}
	return d[m*(width)+0]
}

func getPrefix(query string) string {
	qLen := Min(len(query), 4)
	q := query[0:qLen]
	return strings.ToLower(q)
}

type Document struct {
	docId int
	bloom int
}

//Used for the bloom filter
const (
	FNV_BASIS_64 = uint64(14695981039346656037)
	FNV_PRIME_64 = uint64((1 << 40) + 435)
	FNV_MASK_64  = uint64(^uint64(0) >> 1)
	NUM_BITS     = 64

	FNV_BASIS_32 = uint32(0x811c9dc5)
	FNV_PRIME_32 = uint32((1 << 24) + 403)
	FNV_MASK_32  = uint32(^uint32(0) >> 1)
)

//The bloom filter of a word is 8 bytes in length
//and has each character added separately
func computeBloomFilter(s string) int {
	cnt := len(s)

	if cnt <= 0 {
		return 0
	}

	var filter int
	hash := uint64(0)

	for i := 0; i < cnt; i++ {
		c := s[i]

		//first hash function
		hash ^= uint64(0xFF & c)
		hash *= FNV_PRIME_64

		//second hash function (reduces collisions for bloom)
		hash ^= uint64(0xFF & (c >> 16))
		hash *= FNV_PRIME_64

		//position of the bit mod the number of bits (8 bytes = 64 bits)
		bitpos := hash % NUM_BITS
		if bitpos < 0 {
			bitpos += NUM_BITS
		}
		filter = filter | (1 << bitpos)
	}

	return filter
}

//Inverted Index - Maps the query prefix to the matching documents
type InvertedIndex map[string][]Document

func NewInvertedIndex() *InvertedIndex {
	i := make(InvertedIndex)
	return &i
}

func (x *InvertedIndex) Size() int {
	return len(map[string][]Document(*x))
}

func (x *InvertedIndex) AddDoc(docId int, doc string, bloom int) {
	for _, word := range strings.Fields(doc) {
		word = getPrefix(word)

		ref, ok := (*x)[word]
		if !ok {
			ref = nil
		}

		(*x)[word] = append(ref, Document{docId: docId, bloom: bloom})
	}
}

func (x *InvertedIndex) Search(query string) []Document {
	q := getPrefix(query)

	ref, ok := (*x)[q]

	if ok {
		return ref
	}
	return nil
}

//Forward Index - Maps the document id to the document
type ForwardIndex map[int]string

func NewForwardIndex() *ForwardIndex {
	i := make(ForwardIndex)
	return &i
}
func (x *ForwardIndex) AddDoc(docId int, doc string) {
	for _, word := range strings.Fields(doc) {
		_, ok := (*x)[docId]
		if !ok {
			(*x)[docId] = word
		}
	}
}
func (x *ForwardIndex) itemAt(i int) string {
	return (*x)[i]
}
<|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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
//
//     http://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 oglemock_test

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"github.com/jacobsa/oglemock"
	"reflect"
	"testing"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type ReturnTest struct {
}

func init()                     { RegisterTestSuite(&ReturnTest{}) }
func TestOgletest(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func (t *ReturnTest) NoReturnValues() {
	sig := reflect.TypeOf(func() {})
	var a oglemock.Action
	var err error
	var vals []interface{}

	// No values.
	a = oglemock.Return()
	err = a.CheckType(sig)
	ExpectEq(nil, err)

	vals = a.Invoke([]interface{}{})
	ExpectThat(vals, ElementsAre())

	// One value.
	a = oglemock.Return(17)
	err = a.CheckType(sig)
	ExpectThat(err, HasSubstr("given 1 val"))
	ExpectThat(err, HasSubstr("expected 0"))

	// Two values.
	a = oglemock.Return(17, 19)
	err = a.CheckType(sig)
	ExpectThat(err, HasSubstr("given 2 vals"))
	ExpectThat(err, HasSubstr("expected 0"))
}

func (t *ReturnTest) Bool() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int8() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int16() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int32() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint8() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint16() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint32() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uintptr() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Float32() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Float64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Complex64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Complex128() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) ArrayOfInt() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) ChanOfInt() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Func() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Interface() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) MapFromStringToInt() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) PointerToString() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) SliceOfInts() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) String() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Struct() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) UnsafePointer() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) MultipleReturnValues() {
	ExpectTrue(false, "TODO")
}
<commit_msg>ReturnTest.Bool<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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
//
//     http://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 oglemock_test

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"github.com/jacobsa/oglemock"
	"reflect"
	"testing"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type ReturnTest struct {
}

func init()                     { RegisterTestSuite(&ReturnTest{}) }
func TestOgletest(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func (t *ReturnTest) NoReturnValues() {
	sig := reflect.TypeOf(func() {})
	var a oglemock.Action
	var err error
	var vals []interface{}

	// No values.
	a = oglemock.Return()
	err = a.CheckType(sig)
	AssertEq(nil, err)

	vals = a.Invoke([]interface{}{})
	ExpectThat(vals, ElementsAre())

	// One value.
	a = oglemock.Return(17)
	err = a.CheckType(sig)
	ExpectThat(err, Error(HasSubstr("given 1 val")))
	ExpectThat(err, Error(HasSubstr("expected 0")))

	// Two values.
	a = oglemock.Return(17, 19)
	err = a.CheckType(sig)
	ExpectThat(err, Error(HasSubstr("given 2 vals")))
	ExpectThat(err, Error(HasSubstr("expected 0")))
}

func (t *ReturnTest) Bool() {
	sig := reflect.TypeOf(func() bool { return false })
	var a oglemock.Action
	var err error
	var vals []interface{}

	// True
	a = oglemock.Return(true)
	err = a.CheckType(sig)
	AssertEq(nil, err)

	vals = a.Invoke([]interface{}{})
	ExpectThat(vals, ElementsAre(true))

	// False
	a = oglemock.Return(false)
	err = a.CheckType(sig)
	AssertEq(nil, err)

	vals = a.Invoke([]interface{}{})
	ExpectThat(vals, ElementsAre(false))

	// Int value
	a = oglemock.Return(int(17))
	err = a.CheckType(sig)
	ExpectThat(err, Error(HasSubstr("given int")))
	ExpectThat(err, Error(HasSubstr("expected bool")))

	// String value
	a = oglemock.Return("taco")
	err = a.CheckType(sig)
	ExpectThat(err, Error(HasSubstr("given string")))
	ExpectThat(err, Error(HasSubstr("expected bool")))
}

func (t *ReturnTest) Int() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int8() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int16() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int32() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Int64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint8() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint16() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint32() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uint64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Uintptr() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Float32() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Float64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Complex64() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Complex128() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) ArrayOfInt() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) ChanOfInt() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Func() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Interface() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) MapFromStringToInt() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) PointerToString() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) SliceOfInts() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) String() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) Struct() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) UnsafePointer() {
	ExpectTrue(false, "TODO")
}

func (t *ReturnTest) MultipleReturnValues() {
	ExpectTrue(false, "TODO")
}
<|endoftext|>


<commit_before>package route

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/eBay/fabio/metrics"
)

var errInvalidPrefix = errors.New("route: prefix must not be empty")
var errInvalidTarget = errors.New("route: target must not be empty")
var errNoMatch = errors.New("route: no target match")

// table stores the active routing table. Must never be nil.
var table atomic.Value

// init initializes the routing table.
func init() {
	table.Store(make(Table))
}

// GetTable returns the active routing table. The function
// is safe to be called from multiple goroutines and the
// value is never nil.
func GetTable() Table {
	return table.Load().(Table)
}

// mu guards table and registry in SetTable.
var mu sync.Mutex

// SetTable sets the active routing table. A nil value
// logs a warning and is ignored. The function is safe
// to be called from multiple goroutines.
func SetTable(t Table) {
	if t == nil {
		log.Print("[WARN] Ignoring nil routing table")
		return
	}
	mu.Lock()
	table.Store(t)
	syncRegistry(t)
	mu.Unlock()
	log.Printf("[INFO] Updated config to\n%s", t)
}

// syncRegistry unregisters all inactive timers.
// It assumes that all timers of the table have
// already been registered.
func syncRegistry(t Table) {
	timers := map[string]bool{}

	// get all registered timers
	metrics.ServiceRegistry.Each(func(name string, m interface{}) {
		timers[name] = false
	})

	// mark the ones from this table as active.
	// this can also add new entries but we do not
	// really care since we are only interested in the
	// inactive ones.
	for _, routes := range t {
		for _, r := range routes {
			for _, tg := range r.Targets {
				timers[tg.timerName] = true
			}
		}
	}

	// unregister inactive timers
	for name, active := range timers {
		if !active {
			metrics.ServiceRegistry.Unregister(name)
			log.Printf("[INFO] Unregistered timer %s", name)
		}
	}
}

// Table contains a set of routes grouped by host.
// The host routes are sorted from most to least specific
// by sorting the routes in reverse order by path.
type Table map[string]Routes

// hostpath splits a host/path prefix into a host and a path.
// The path always starts with a slash
func hostpath(prefix string) (host string, path string) {
	p := strings.SplitN(prefix, "/", 2)
	host, path = p[0], ""
	if len(p) == 1 {
		return p[0], "/"
	}
	return p[0], "/" + p[1]
}

// AddRoute adds a new route prefix -> target for the given service.
func (t Table) AddRoute(service, prefix, target string, weight float64, tags []string) error {
	host, path := hostpath(prefix)

	if prefix == "" {
		return errInvalidPrefix
	}

	if target == "" {
		return errInvalidTarget
	}

	targetURL, err := url.Parse(target)
	if err != nil {
		return fmt.Errorf("route: invalid target. %s", err)
	}

	r := newRoute(host, path)
	r.addTarget(service, targetURL, weight, tags)

	// add new host
	if t[host] == nil {
		t[host] = Routes{r}
		return nil
	}

	// add new route to existing host
	if t[host].find(path) == nil {
		t[host] = append(t[host], r)
		sort.Sort(t[host])
		return nil
	}

	// add new target to existing route
	t[host].find(path).addTarget(service, targetURL, weight, tags)

	return nil
}

func (t Table) AddRouteWeight(service, prefix string, weight float64, tags []string) error {
	host, path := hostpath(prefix)

	if prefix == "" {
		return errInvalidPrefix
	}

	if t[host] == nil || t[host].find(path) == nil {
		return errNoMatch
	}

	if n := t[host].find(path).setWeight(service, weight, tags); n == 0 {
		return errNoMatch
	}
	return nil
}

// DelRoute removes one or more routes depending on the arguments.
// If service, prefix and target are provided then only this route
// is removed. Are only service and prefix provided then all routes
// for this service and prefix are removed. This removes all active
// instances of the service from the route. If only the service is
// provided then all routes for this service are removed. The service
// will no longer receive traffic.
func (t Table) DelRoute(service, prefix, target string) error {
	switch {
	case prefix == "" && target == "":
		for _, hr := range t {
			for _, r := range hr {
				r.delService(service)
			}
		}

	case target == "":
		r := t.route(hostpath(prefix))
		if r == nil {
			return nil
		}
		r.delService(service)

	default:
		targetURL, err := url.Parse(target)
		if err != nil {
			return fmt.Errorf("route: invalid target. %s", err)
		}

		r := t.route(hostpath(prefix))
		if r == nil {
			return nil
		}
		r.delTarget(service, targetURL)
	}

	return nil
}

// route finds the route for host/path or returns nil if none exists.
func (t Table) route(host, path string) *Route {
	hr := t[host]
	if hr == nil {
		return nil
	}
	return hr.find(path)
}

// Lookup finds a target url based on the current matcher and picker
// or nil if there is none. It first checks the routes for the host
// and if none matches then it falls back to generic routes without
// a host. This is useful for a catch-all '/' rule.
func (t Table) Lookup(req *http.Request, trace string) *Target {
	if trace != "" {
		if len(trace) > 16 {
			trace = trace[:15]
		}
		log.Printf("[TRACE] %s Tracing %s%s", trace, req.Host, req.RequestURI)
	}

	u := t.doLookup(strings.ToLower(req.Host), req.RequestURI, trace)
	if u == nil {
		u = t.doLookup("", req.RequestURI, trace)
	}

	if trace != "" {
		log.Printf("[TRACE] %s Routing to %s", trace, u.URL)
	}

	return u
}

func (t Table) doLookup(host, path, trace string) *Target {
	hr := t[host]
	if hr == nil {
		return nil
	}

	for _, r := range hr {
		if match(path, r) {
			n := len(r.Targets)
			if n == 0 {
				return nil
			}
			if n == 1 {
				return r.Targets[0]
			}
			if trace != "" {
				log.Printf("[TRACE] %s Match %s%s", trace, r.Host, r.Path)
			}
			return pick(r)
		}
		if trace != "" {
			log.Printf("[TRACE] %s No match %s%s", trace, r.Host, r.Path)
		}
	}
	return nil
}

func (t Table) Config(addWeight bool) []string {
	var hosts []string
	for h := range t {
		if h != "" {
			hosts = append(hosts, h)
		}
	}
	sort.Strings(hosts)

	// entries without host come last
	hosts = append(hosts, "")

	var routes []string
	for _, h := range hosts {
		for _, r := range t[h] {
			routes = append(routes, r.config(addWeight)...)
		}
	}
	return routes
}

// String returns the routing table as config file which can
// be read by Parse() again.
func (t Table) String() string {
	return strings.Join(t.Config(false), "\n")
}
<commit_msg>Rename variables<commit_after>package route

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/eBay/fabio/metrics"
)

var errInvalidPrefix = errors.New("route: prefix must not be empty")
var errInvalidTarget = errors.New("route: target must not be empty")
var errNoMatch = errors.New("route: no target match")

// table stores the active routing table. Must never be nil.
var table atomic.Value

// init initializes the routing table.
func init() {
	table.Store(make(Table))
}

// GetTable returns the active routing table. The function
// is safe to be called from multiple goroutines and the
// value is never nil.
func GetTable() Table {
	return table.Load().(Table)
}

// mu guards table and registry in SetTable.
var mu sync.Mutex

// SetTable sets the active routing table. A nil value
// logs a warning and is ignored. The function is safe
// to be called from multiple goroutines.
func SetTable(t Table) {
	if t == nil {
		log.Print("[WARN] Ignoring nil routing table")
		return
	}
	mu.Lock()
	table.Store(t)
	syncRegistry(t)
	mu.Unlock()
	log.Printf("[INFO] Updated config to\n%s", t)
}

// syncRegistry unregisters all inactive timers.
// It assumes that all timers of the table have
// already been registered.
func syncRegistry(t Table) {
	timers := map[string]bool{}

	// get all registered timers
	metrics.ServiceRegistry.Each(func(name string, m interface{}) {
		timers[name] = false
	})

	// mark the ones from this table as active.
	// this can also add new entries but we do not
	// really care since we are only interested in the
	// inactive ones.
	for _, routes := range t {
		for _, r := range routes {
			for _, tg := range r.Targets {
				timers[tg.timerName] = true
			}
		}
	}

	// unregister inactive timers
	for name, active := range timers {
		if !active {
			metrics.ServiceRegistry.Unregister(name)
			log.Printf("[INFO] Unregistered timer %s", name)
		}
	}
}

// Table contains a set of routes grouped by host.
// The host routes are sorted from most to least specific
// by sorting the routes in reverse order by path.
type Table map[string]Routes

// hostpath splits a host/path prefix into a host and a path.
// The path always starts with a slash
func hostpath(prefix string) (host string, path string) {
	p := strings.SplitN(prefix, "/", 2)
	host, path = p[0], ""
	if len(p) == 1 {
		return p[0], "/"
	}
	return p[0], "/" + p[1]
}

// AddRoute adds a new route prefix -> target for the given service.
func (t Table) AddRoute(service, prefix, target string, weight float64, tags []string) error {
	host, path := hostpath(prefix)

	if prefix == "" {
		return errInvalidPrefix
	}

	if target == "" {
		return errInvalidTarget
	}

	targetURL, err := url.Parse(target)
	if err != nil {
		return fmt.Errorf("route: invalid target. %s", err)
	}

	r := newRoute(host, path)
	r.addTarget(service, targetURL, weight, tags)

	// add new host
	if t[host] == nil {
		t[host] = Routes{r}
		return nil
	}

	// add new route to existing host
	if t[host].find(path) == nil {
		t[host] = append(t[host], r)
		sort.Sort(t[host])
		return nil
	}

	// add new target to existing route
	t[host].find(path).addTarget(service, targetURL, weight, tags)

	return nil
}

func (t Table) AddRouteWeight(service, prefix string, weight float64, tags []string) error {
	host, path := hostpath(prefix)

	if prefix == "" {
		return errInvalidPrefix
	}

	if t[host] == nil || t[host].find(path) == nil {
		return errNoMatch
	}

	if n := t[host].find(path).setWeight(service, weight, tags); n == 0 {
		return errNoMatch
	}
	return nil
}

// DelRoute removes one or more routes depending on the arguments.
// If service, prefix and target are provided then only this route
// is removed. Are only service and prefix provided then all routes
// for this service and prefix are removed. This removes all active
// instances of the service from the route. If only the service is
// provided then all routes for this service are removed. The service
// will no longer receive traffic.
func (t Table) DelRoute(service, prefix, target string) error {
	switch {
	case prefix == "" && target == "":
		for _, routes := range t {
			for _, r := range routes {
				r.delService(service)
			}
		}

	case target == "":
		r := t.route(hostpath(prefix))
		if r == nil {
			return nil
		}
		r.delService(service)

	default:
		targetURL, err := url.Parse(target)
		if err != nil {
			return fmt.Errorf("route: invalid target. %s", err)
		}

		r := t.route(hostpath(prefix))
		if r == nil {
			return nil
		}
		r.delTarget(service, targetURL)
	}

	return nil
}

// route finds the route for host/path or returns nil if none exists.
func (t Table) route(host, path string) *Route {
	routes := t[host]
	if routes == nil {
		return nil
	}
	return routes.find(path)
}

// Lookup finds a target url based on the current matcher and picker
// or nil if there is none. It first checks the routes for the host
// and if none matches then it falls back to generic routes without
// a host. This is useful for a catch-all '/' rule.
func (t Table) Lookup(req *http.Request, trace string) *Target {
	if trace != "" {
		if len(trace) > 16 {
			trace = trace[:15]
		}
		log.Printf("[TRACE] %s Tracing %s%s", trace, req.Host, req.RequestURI)
	}

	target := t.doLookup(strings.ToLower(req.Host), req.RequestURI, trace)
	if target == nil {
		target = t.doLookup("", req.RequestURI, trace)
	}

	if trace != "" {
		log.Printf("[TRACE] %s Routing to %s", trace, target.URL)
	}

	return target
}

func (t Table) doLookup(host, path, trace string) *Target {
	routes := t[host]
	if routes == nil {
		return nil
	}

	for _, r := range routes {
		if match(path, r) {
			n := len(r.Targets)
			if n == 0 {
				return nil
			}
			if n == 1 {
				return r.Targets[0]
			}
			if trace != "" {
				log.Printf("[TRACE] %s Match %s%s", trace, r.Host, r.Path)
			}
			return pick(r)
		}
		if trace != "" {
			log.Printf("[TRACE] %s No match %s%s", trace, r.Host, r.Path)
		}
	}
	return nil
}

func (t Table) Config(addWeight bool) []string {
	var hosts []string
	for host := range t {
		if host != "" {
			hosts = append(hosts, host)
		}
	}
	sort.Strings(hosts)

	// entries without host come last
	hosts = append(hosts, "")

	var cfg []string
	for _, host := range hosts {
		for _, routes := range t[host] {
			cfg = append(cfg, routes.config(addWeight)...)
		}
	}
	return cfg
}

// String returns the routing table as config file which can
// be read by Parse() again.
func (t Table) String() string {
	return strings.Join(t.Config(false), "\n")
}
<|endoftext|>


<commit_before>package netlink

import (
	"fmt"
	"net"
	"syscall"

	"github.com/vishvananda/netlink/nl"
)

// RtAttr is shared so it is in netlink_linux.go

// RouteAdd will add a route to the system.
// Equivalent to: `ip route add $route`
func RouteAdd(route *Route) error {
	req := nl.NewNetlinkRequest(syscall.RTM_NEWROUTE, syscall.NLM_F_CREATE|syscall.NLM_F_EXCL|syscall.NLM_F_ACK)
	return routeHandle(route, req, nl.NewRtMsg())
}

// RouteAdd will delete a route from the system.
// Equivalent to: `ip route del $route`
func RouteDel(route *Route) error {
	req := nl.NewNetlinkRequest(syscall.RTM_DELROUTE, syscall.NLM_F_ACK)
	return routeHandle(route, req, nl.NewRtDelMsg())
}

func routeHandle(route *Route, req *nl.NetlinkRequest, msg *nl.RtMsg) error {
	if (route.Dst == nil || route.Dst.IP == nil) && route.Src == nil && route.Gw == nil {
		return fmt.Errorf("one of Dst.IP, Src, or Gw must not be nil")
	}

	msg.Scope = uint8(route.Scope)
	family := -1
	var rtAttrs []*nl.RtAttr

	if route.Dst != nil && route.Dst.IP != nil {
		dstLen, _ := route.Dst.Mask.Size()
		msg.Dst_len = uint8(dstLen)
		dstFamily := nl.GetIPFamily(route.Dst.IP)
		family = dstFamily
		var dstData []byte
		if dstFamily == FAMILY_V4 {
			dstData = route.Dst.IP.To4()
		} else {
			dstData = route.Dst.IP.To16()
		}
		rtAttrs = append(rtAttrs, nl.NewRtAttr(syscall.RTA_DST, dstData))
	}

	if route.Src != nil {
		srcFamily := nl.GetIPFamily(route.Src)
		if family != -1 && family != srcFamily {
			return fmt.Errorf("source and destination ip are not the same IP family")
		}
		family = srcFamily
		var srcData []byte
		if srcFamily == FAMILY_V4 {
			srcData = route.Src.To4()
		} else {
			srcData = route.Src.To16()
		}
		// The commonly used src ip for routes is actually PREFSRC
		rtAttrs = append(rtAttrs, nl.NewRtAttr(syscall.RTA_PREFSRC, srcData))
	}

	if route.Gw != nil {
		gwFamily := nl.GetIPFamily(route.Gw)
		if family != -1 && family != gwFamily {
			return fmt.Errorf("gateway, source, and destination ip are not the same IP family")
		}
		family = gwFamily
		var gwData []byte
		if gwFamily == FAMILY_V4 {
			gwData = route.Gw.To4()
		} else {
			gwData = route.Gw.To16()
		}
		rtAttrs = append(rtAttrs, nl.NewRtAttr(syscall.RTA_GATEWAY, gwData))
	}

	msg.Family = uint8(family)

	req.AddData(msg)
	for _, attr := range rtAttrs {
		req.AddData(attr)
	}

	var (
		b      = make([]byte, 4)
		native = nl.NativeEndian()
	)
	native.PutUint32(b, uint32(route.LinkIndex))

	req.AddData(nl.NewRtAttr(syscall.RTA_OIF, b))

	_, err := req.Execute(syscall.NETLINK_ROUTE, 0)
	return err
}

// RouteList gets a list of routes in the system.
// Equivalent to: `ip route show`.
// The list can be filtered by link and ip family.
func RouteList(link Link, family int) ([]Route, error) {
	req := nl.NewNetlinkRequest(syscall.RTM_GETROUTE, syscall.NLM_F_DUMP)
	msg := nl.NewIfInfomsg(family)
	req.AddData(msg)

	msgs, err := req.Execute(syscall.NETLINK_ROUTE, syscall.RTM_NEWROUTE)
	if err != nil {
		return nil, err
	}

	index := 0
	if link != nil {
		base := link.Attrs()
		ensureIndex(base)
		index = base.Index
	}

	native := nl.NativeEndian()
	var res []Route
MsgLoop:
	for _, m := range msgs {
		msg := nl.DeserializeRtMsg(m)

		if msg.Flags&syscall.RTM_F_CLONED != 0 {
			// Ignore cloned routes
			continue
		}

		if msg.Table != syscall.RT_TABLE_MAIN {
			// Ignore non-main tables
			continue
		}

		attrs, err := nl.ParseRouteAttr(m[msg.Len():])
		if err != nil {
			return nil, err
		}

		route := Route{Scope: Scope(msg.Scope)}
		for _, attr := range attrs {
			switch attr.Attr.Type {
			case syscall.RTA_GATEWAY:
				route.Gw = net.IP(attr.Value)
			case syscall.RTA_PREFSRC:
				route.Src = net.IP(attr.Value)
			case syscall.RTA_DST:
				route.Dst = &net.IPNet{
					IP:   attr.Value,
					Mask: net.CIDRMask(int(msg.Dst_len), 8*len(attr.Value)),
				}
			case syscall.RTA_OIF:
				routeIndex := int(native.Uint32(attr.Value[0:4]))
				if link != nil && routeIndex != index {
					// Ignore routes from other interfaces
					continue MsgLoop
				}
				route.LinkIndex = routeIndex
			}
		}
		res = append(res, route)
	}

	return res, nil
}

// RouteGet gets a route to a specific destination from the host system.
// Equivalent to: 'ip route get'.
func RouteGet(destination net.IP) ([]Route, error) {
	req := nl.NewNetlinkRequest(syscall.RTM_GETROUTE, syscall.NLM_F_REQUEST)
	family := nl.GetIPFamily(destination)
	var destinationData []byte
	var bitlen uint8
	if family == FAMILY_V4 {
		destinationData = destination.To4()
		bitlen = 32
	} else {
		destinationData = destination.To16()
		bitlen = 128
	}
	msg := &nl.RtMsg{}
	msg.Family = uint8(family)
	msg.Dst_len = bitlen
	req.AddData(msg)

	rtaDst := nl.NewRtAttr(syscall.RTA_DST, destinationData)
	req.AddData(rtaDst)

	msgs, err := req.Execute(syscall.NETLINK_ROUTE, syscall.RTM_NEWROUTE)
	if err != nil {
		return nil, err
	}

	native := nl.NativeEndian()
	var res []Route
	for _, m := range msgs {
		msg := nl.DeserializeRtMsg(m)
		attrs, err := nl.ParseRouteAttr(m[msg.Len():])
		if err != nil {
			return nil, err
		}

		route := Route{}
		for _, attr := range attrs {
			switch attr.Attr.Type {
			case syscall.RTA_GATEWAY:
				route.Gw = net.IP(attr.Value)
			case syscall.RTA_PREFSRC:
				route.Src = net.IP(attr.Value)
			case syscall.RTA_DST:
				route.Dst = &net.IPNet{
					IP:   attr.Value,
					Mask: net.CIDRMask(int(msg.Dst_len), 8*len(attr.Value)),
				}
			case syscall.RTA_OIF:
				routeIndex := int(native.Uint32(attr.Value[0:4]))
				route.LinkIndex = routeIndex
			}
		}
		res = append(res, route)
	}
	return res, nil

}
<commit_msg>Fix typo<commit_after>package netlink

import (
	"fmt"
	"net"
	"syscall"

	"github.com/vishvananda/netlink/nl"
)

// RtAttr is shared so it is in netlink_linux.go

// RouteAdd will add a route to the system.
// Equivalent to: `ip route add $route`
func RouteAdd(route *Route) error {
	req := nl.NewNetlinkRequest(syscall.RTM_NEWROUTE, syscall.NLM_F_CREATE|syscall.NLM_F_EXCL|syscall.NLM_F_ACK)
	return routeHandle(route, req, nl.NewRtMsg())
}

// RouteDel will delete a route from the system.
// Equivalent to: `ip route del $route`
func RouteDel(route *Route) error {
	req := nl.NewNetlinkRequest(syscall.RTM_DELROUTE, syscall.NLM_F_ACK)
	return routeHandle(route, req, nl.NewRtDelMsg())
}

func routeHandle(route *Route, req *nl.NetlinkRequest, msg *nl.RtMsg) error {
	if (route.Dst == nil || route.Dst.IP == nil) && route.Src == nil && route.Gw == nil {
		return fmt.Errorf("one of Dst.IP, Src, or Gw must not be nil")
	}

	msg.Scope = uint8(route.Scope)
	family := -1
	var rtAttrs []*nl.RtAttr

	if route.Dst != nil && route.Dst.IP != nil {
		dstLen, _ := route.Dst.Mask.Size()
		msg.Dst_len = uint8(dstLen)
		dstFamily := nl.GetIPFamily(route.Dst.IP)
		family = dstFamily
		var dstData []byte
		if dstFamily == FAMILY_V4 {
			dstData = route.Dst.IP.To4()
		} else {
			dstData = route.Dst.IP.To16()
		}
		rtAttrs = append(rtAttrs, nl.NewRtAttr(syscall.RTA_DST, dstData))
	}

	if route.Src != nil {
		srcFamily := nl.GetIPFamily(route.Src)
		if family != -1 && family != srcFamily {
			return fmt.Errorf("source and destination ip are not the same IP family")
		}
		family = srcFamily
		var srcData []byte
		if srcFamily == FAMILY_V4 {
			srcData = route.Src.To4()
		} else {
			srcData = route.Src.To16()
		}
		// The commonly used src ip for routes is actually PREFSRC
		rtAttrs = append(rtAttrs, nl.NewRtAttr(syscall.RTA_PREFSRC, srcData))
	}

	if route.Gw != nil {
		gwFamily := nl.GetIPFamily(route.Gw)
		if family != -1 && family != gwFamily {
			return fmt.Errorf("gateway, source, and destination ip are not the same IP family")
		}
		family = gwFamily
		var gwData []byte
		if gwFamily == FAMILY_V4 {
			gwData = route.Gw.To4()
		} else {
			gwData = route.Gw.To16()
		}
		rtAttrs = append(rtAttrs, nl.NewRtAttr(syscall.RTA_GATEWAY, gwData))
	}

	msg.Family = uint8(family)

	req.AddData(msg)
	for _, attr := range rtAttrs {
		req.AddData(attr)
	}

	var (
		b      = make([]byte, 4)
		native = nl.NativeEndian()
	)
	native.PutUint32(b, uint32(route.LinkIndex))

	req.AddData(nl.NewRtAttr(syscall.RTA_OIF, b))

	_, err := req.Execute(syscall.NETLINK_ROUTE, 0)
	return err
}

// RouteList gets a list of routes in the system.
// Equivalent to: `ip route show`.
// The list can be filtered by link and ip family.
func RouteList(link Link, family int) ([]Route, error) {
	req := nl.NewNetlinkRequest(syscall.RTM_GETROUTE, syscall.NLM_F_DUMP)
	msg := nl.NewIfInfomsg(family)
	req.AddData(msg)

	msgs, err := req.Execute(syscall.NETLINK_ROUTE, syscall.RTM_NEWROUTE)
	if err != nil {
		return nil, err
	}

	index := 0
	if link != nil {
		base := link.Attrs()
		ensureIndex(base)
		index = base.Index
	}

	native := nl.NativeEndian()
	var res []Route
MsgLoop:
	for _, m := range msgs {
		msg := nl.DeserializeRtMsg(m)

		if msg.Flags&syscall.RTM_F_CLONED != 0 {
			// Ignore cloned routes
			continue
		}

		if msg.Table != syscall.RT_TABLE_MAIN {
			// Ignore non-main tables
			continue
		}

		attrs, err := nl.ParseRouteAttr(m[msg.Len():])
		if err != nil {
			return nil, err
		}

		route := Route{Scope: Scope(msg.Scope)}
		for _, attr := range attrs {
			switch attr.Attr.Type {
			case syscall.RTA_GATEWAY:
				route.Gw = net.IP(attr.Value)
			case syscall.RTA_PREFSRC:
				route.Src = net.IP(attr.Value)
			case syscall.RTA_DST:
				route.Dst = &net.IPNet{
					IP:   attr.Value,
					Mask: net.CIDRMask(int(msg.Dst_len), 8*len(attr.Value)),
				}
			case syscall.RTA_OIF:
				routeIndex := int(native.Uint32(attr.Value[0:4]))
				if link != nil && routeIndex != index {
					// Ignore routes from other interfaces
					continue MsgLoop
				}
				route.LinkIndex = routeIndex
			}
		}
		res = append(res, route)
	}

	return res, nil
}

// RouteGet gets a route to a specific destination from the host system.
// Equivalent to: 'ip route get'.
func RouteGet(destination net.IP) ([]Route, error) {
	req := nl.NewNetlinkRequest(syscall.RTM_GETROUTE, syscall.NLM_F_REQUEST)
	family := nl.GetIPFamily(destination)
	var destinationData []byte
	var bitlen uint8
	if family == FAMILY_V4 {
		destinationData = destination.To4()
		bitlen = 32
	} else {
		destinationData = destination.To16()
		bitlen = 128
	}
	msg := &nl.RtMsg{}
	msg.Family = uint8(family)
	msg.Dst_len = bitlen
	req.AddData(msg)

	rtaDst := nl.NewRtAttr(syscall.RTA_DST, destinationData)
	req.AddData(rtaDst)

	msgs, err := req.Execute(syscall.NETLINK_ROUTE, syscall.RTM_NEWROUTE)
	if err != nil {
		return nil, err
	}

	native := nl.NativeEndian()
	var res []Route
	for _, m := range msgs {
		msg := nl.DeserializeRtMsg(m)
		attrs, err := nl.ParseRouteAttr(m[msg.Len():])
		if err != nil {
			return nil, err
		}

		route := Route{}
		for _, attr := range attrs {
			switch attr.Attr.Type {
			case syscall.RTA_GATEWAY:
				route.Gw = net.IP(attr.Value)
			case syscall.RTA_PREFSRC:
				route.Src = net.IP(attr.Value)
			case syscall.RTA_DST:
				route.Dst = &net.IPNet{
					IP:   attr.Value,
					Mask: net.CIDRMask(int(msg.Dst_len), 8*len(attr.Value)),
				}
			case syscall.RTA_OIF:
				routeIndex := int(native.Uint32(attr.Value[0:4]))
				route.LinkIndex = routeIndex
			}
		}
		res = append(res, route)
	}
	return res, nil

}
<|endoftext|>


<commit_before>package router

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

var benchmarkRouter *Router

func getBenchmarkRouter() *Router {
	if benchmarkRouter != nil {
		return benchmarkRouter
	}

	benchmarkRouter := NewRouter()
	handler := func(http.ResponseWriter, *http.Request) {}
	benchmarkRouter.AddRoute("/", handler)
	benchmarkRouter.AddRoute("/foo", handler)
	benchmarkRouter.AddRoute("/foo/bar", handler)
	benchmarkRouter.AddRoute("/foo/baz", handler)
	benchmarkRouter.AddRoute("/people", handler)
	benchmarkRouter.AddRoute("/people/search", handler)
	benchmarkRouter.AddRoute("/people/:id", handler)
	benchmarkRouter.AddRoute("/users", handler)
	benchmarkRouter.AddRoute("/users/:id", handler)
	benchmarkRouter.AddRoute("/widgets", handler)
	benchmarkRouter.AddRoute("/widgets/important", handler)

	return benchmarkRouter
}

func TestSegmentizePath(t *testing.T) {
	test := func(path string, expected []string) {
		actual := segmentizePath(path)
		if len(actual) != len(expected) {
			t.Errorf("Expected \"%v\" to be segmented into %v, but it actually was %v", path, expected, actual)
			return
		}

		for i := 0; i < len(actual); i++ {
			if actual[i] != expected[i] {
				t.Errorf("Expected \"%v\" to be segmented into %v, but it actually was %v", path, expected, actual)
				return
			}
		}
	}

	test("/", []string{})
	test("/foo", []string{"foo"})
	test("/foo/", []string{"foo"})
	test("/foo/bar", []string{"foo", "bar"})
	test("/foo/bar/", []string{"foo", "bar"})
	test("/foo/bar/baz", []string{"foo", "bar", "baz"})
}

func TestRouter(t *testing.T) {
	router := NewRouter()

	rootHandler := func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "root")
	}
	router.AddRoute("/", rootHandler)

	widgetIndexHandler := func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "widgetIndex")
	}
	router.AddRoute("/widget", widgetIndexHandler)

	widgetShowHandler := func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "widgetShow")
	}
	router.AddRoute("/widget/?", widgetShowHandler)

	get := func(path string, expectedCode int, expectedBody string) {
		response := httptest.NewRecorder()
		request, err := http.NewRequest("GET", "http://example.com"+path, nil)
		if err != nil {
			t.Errorf("Unable to create test GET request for %v", path)
		}

		router.ServeHTTP(response, request)
		if response.Code != expectedCode {
			t.Errorf("GET %v: expected HTTP code %v, received %v", path, expectedCode, response.Code)
		}
		if response.Body.String() != expectedBody {
			t.Errorf("GET %v: expected HTTP response body \"%v\", received \"%v\"", path, expectedBody, response.Body.String())
		}
	}

	get("/", 200, "root")
	get("/widget", 200, "widgetIndex")
	get("/widget/1", 200, "widgetShow")

	get("/missing", 404, "404 Not Found")
	get("/widget/1/missing", 404, "404 Not Found")
}

func BenchmarkFindHandlerRoot(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/"))
	}
}

func BenchmarkFindHandlerSingleLevel(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/foo"))
	}
}

func BenchmarkFindHandlerSecondLevel(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/people/search"))
	}
}
<commit_msg>Add more benchmarks<commit_after>package router

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

var benchmarkRouter *Router

func getBenchmarkRouter() *Router {
	if benchmarkRouter != nil {
		return benchmarkRouter
	}

	benchmarkRouter := NewRouter()
	handler := func(http.ResponseWriter, *http.Request) {}
	benchmarkRouter.AddRoute("/", handler)
	benchmarkRouter.AddRoute("/foo", handler)
	benchmarkRouter.AddRoute("/foo/bar", handler)
	benchmarkRouter.AddRoute("/foo/baz", handler)
	benchmarkRouter.AddRoute("/foo/bar/baz/quz", handler)
	benchmarkRouter.AddRoute("/people", handler)
	benchmarkRouter.AddRoute("/people/search", handler)
	benchmarkRouter.AddRoute("/people/?", handler)
	benchmarkRouter.AddRoute("/users", handler)
	benchmarkRouter.AddRoute("/users/?", handler)
	benchmarkRouter.AddRoute("/widgets", handler)
	benchmarkRouter.AddRoute("/widgets/important", handler)

	return benchmarkRouter
}

func TestSegmentizePath(t *testing.T) {
	test := func(path string, expected []string) {
		actual := segmentizePath(path)
		if len(actual) != len(expected) {
			t.Errorf("Expected \"%v\" to be segmented into %v, but it actually was %v", path, expected, actual)
			return
		}

		for i := 0; i < len(actual); i++ {
			if actual[i] != expected[i] {
				t.Errorf("Expected \"%v\" to be segmented into %v, but it actually was %v", path, expected, actual)
				return
			}
		}
	}

	test("/", []string{})
	test("/foo", []string{"foo"})
	test("/foo/", []string{"foo"})
	test("/foo/bar", []string{"foo", "bar"})
	test("/foo/bar/", []string{"foo", "bar"})
	test("/foo/bar/baz", []string{"foo", "bar", "baz"})
}

func TestRouter(t *testing.T) {
	router := NewRouter()

	rootHandler := func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "root")
	}
	router.AddRoute("/", rootHandler)

	widgetIndexHandler := func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "widgetIndex")
	}
	router.AddRoute("/widget", widgetIndexHandler)

	widgetShowHandler := func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "widgetShow")
	}
	router.AddRoute("/widget/?", widgetShowHandler)

	get := func(path string, expectedCode int, expectedBody string) {
		response := httptest.NewRecorder()
		request, err := http.NewRequest("GET", "http://example.com"+path, nil)
		if err != nil {
			t.Errorf("Unable to create test GET request for %v", path)
		}

		router.ServeHTTP(response, request)
		if response.Code != expectedCode {
			t.Errorf("GET %v: expected HTTP code %v, received %v", path, expectedCode, response.Code)
		}
		if response.Body.String() != expectedBody {
			t.Errorf("GET %v: expected HTTP response body \"%v\", received \"%v\"", path, expectedBody, response.Body.String())
		}
	}

	get("/", 200, "root")
	get("/widget", 200, "widgetIndex")
	get("/widget/1", 200, "widgetShow")

	get("/missing", 404, "404 Not Found")
	get("/widget/1/missing", 404, "404 Not Found")
}

func BenchmarkFindHandlerRoot(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/"))
	}
}

func BenchmarkFindHandlerSegment1(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/foo"))
	}
}

func BenchmarkFindHandlerSegment2(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/people/search"))
	}
}

func BenchmarkFindHandlerSegment2Placeholder(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/people/1"))
	}
}

func BenchmarkFindHandlerSegment4(b *testing.B) {
	router := getBenchmarkRouter()

	for i := 0; i < b.N; i++ {
		router.FindHandler(segmentizePath("/foo/bar/baz/quz"))
	}
}
<|endoftext|>


<commit_before>package styx

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"

	"aqwari.net/net/styx/internal/qidpool"
	"aqwari.net/net/styx/internal/styxfile"
	"aqwari.net/net/styx/internal/tracing"
	"aqwari.net/net/styx/internal/util"
	"aqwari.net/net/styx/styxproto"

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

var (
	errFidInUse     = errors.New("fid already in use")
	errTagInUse     = errors.New("tag in use")
	errNoFid        = errors.New("no such fid")
	errNotSupported = errors.New("not supported")
)

type fcall interface {
	styxproto.Msg
	Fid() uint32
}

// A note on identifiers (fids & tags)
//
// identifiers are chosen by the client, not by the server.  Therefore,
// it is important that the performance and behavior of a server does
// *not* change based on the fid or tag a client chooses. This is why
// a map is used; its performance is good, and doesn't change based
// on the input a client chooses (rather, it does not change in a way
// a client can predict).

// A conn receives and sends 9P messages across a single network connection.
// Multiple "sessions" may take place over a single connection. The conn
// struct contains the necessary information to route 9P messages to their
// established sessions.
type conn struct {
	// These wrap the network connection to read and write messages.
	*styxproto.Decoder
	*styxproto.Encoder

	// The Server a connection was spawned from. Contains configuration
	// settings and the authentication function, if any.
	srv *Server

	// The network connection itself. We expose it in the struct so that
	// it is available for transport-based auth and any timeouts we need
	// to implement.
	rwc io.ReadWriteCloser

	// This serves as the parent context for the context attached to all
	// requests.
	cx context.Context

	// While srv.MaxSize holds the *desired* 9P protocol message
	// size, msize will contain the actual maximum negotiated with
	// the client, through a Tversion/Rversion exchange.
	msize int64

	// There is no "session id" in 9P. However, because all fids
	// for a connection must be derived from the fid established
	// in a Tattach call, any message that contains a fid can be
	// traced back to the original Tattach message.
	sessionFid *util.Map

	// Qids for the file tree, added on-demand.
	qidpool *qidpool.Pool

	// used to implement request cancellation when a Tflush
	// message is received.
	pendingReq map[uint16]context.CancelFunc
}

func (c *conn) remoteAddr() net.Addr {
	type hasRemote interface {
		RemoteAddr() net.Addr
	}
	if nc, ok := c.rwc.(hasRemote); ok {
		return nc.RemoteAddr()
	}
	return nil
}

func (c *conn) sessionByFid(fid uint32) (*Session, bool) {
	if v, ok := c.sessionFid.Get(fid); ok {
		return v.(*Session), true
	}
	return nil, false
}

// Close the connection
func (c *conn) close() error {
	// Cancel all pending requests
	for tag, cancel := range c.pendingReq {
		cancel()
		delete(c.pendingReq, tag)
	}

	// Close all open files and sessions
	c.sessionFid.Do(func(m map[interface{}]interface{}) {
		seen := make(map[*Session]struct{}, len(m))
		for k, v := range m {
			session := v.(*Session)
			if _, ok := seen[session]; !ok {
				seen[session] = struct{}{}
				session.endSession()
			}
			// Should probably let the GC take care of this
			delete(m, k)
		}
	})

	return c.rwc.Close()
}

func newConn(srv *Server, rwc io.ReadWriteCloser) *conn {
	var msize int64 = styxproto.DefaultMaxSize
	if srv.MaxSize > 0 {
		if srv.MaxSize > styxproto.MinBufSize {
			msize = srv.MaxSize
		} else {
			msize = styxproto.MinBufSize
		}
	}
	var enc *styxproto.Encoder
	var dec *styxproto.Decoder
	if srv.TraceLog != nil {
		enc = tracing.Encoder(rwc, func(m styxproto.Msg) {
			srv.TraceLog.Printf("← %03d %s", m.Tag(), m)
		})
		dec = tracing.Decoder(rwc, func(m styxproto.Msg) {
			srv.TraceLog.Printf("→ %03d %s", m.Tag(), m)
		})
	} else {
		enc = styxproto.NewEncoder(rwc)
		dec = styxproto.NewDecoder(rwc)
	}
	return &conn{
		Decoder:    dec,
		Encoder:    enc,
		srv:        srv,
		rwc:        rwc,
		cx:         context.TODO(),
		msize:      msize,
		sessionFid: util.NewMap(),
		pendingReq: make(map[uint16]context.CancelFunc),
		qidpool:    qidpool.New(),
	}
}

func (c *conn) qid(name string, qtype uint8) styxproto.Qid {
	return c.qidpool.Put(name, qtype)
}

// All request contexts must have their cancel functions
// called, to free up resources in the context.
func (c *conn) clearTag(tag uint16) {
	if cancel, ok := c.pendingReq[tag]; ok {
		cancel()
		delete(c.pendingReq, tag)
	}
}

// runs in its own goroutine, one per connection.
func (c *conn) serve() {
	defer c.close()

	if !c.acceptTversion() {
		return
	}

Loop:
	for c.Next() {
		for _, m := range c.Messages() {
			if !c.handleMessage(m) {
				break Loop
			}
		}
	}
	c.srv.logf("closed connection from %s", c.remoteAddr())
}

func (c *conn) handleMessage(m styxproto.Msg) bool {
	if _, ok := c.pendingReq[m.Tag()]; ok {
		c.Rerror(m.Tag(), "%s", errTagInUse)
		return false
	}
	cx, cancel := context.WithCancel(c.cx)
	c.pendingReq[m.Tag()] = cancel

	switch m := m.(type) {
	case styxproto.Tauth:
		return c.handleTauth(cx, m)
	case styxproto.Tattach:
		return c.handleTattach(cx, m)
	case styxproto.Tflush:
		return c.handleTflush(cx, m)
	case fcall:
		return c.handleFcall(cx, m)
	case styxproto.BadMessage:
		c.srv.logf("got bad message from %s: %s", c.remoteAddr(), m.Err)
		c.Rerror(m.Tag(), "bad message: %s", m.Err)
		c.clearTag(m.Tag())
		return true
	default:
		c.Rerror(m.Tag(), "unexpected %T message", m)
		return false
	}
	return true
}

// This is the first thing we do on a new connection. The first
// message a client sends *must* be a Tversion message.
func (c *conn) acceptTversion() bool {
	c.Encoder.MaxSize = c.msize
	c.Decoder.MaxSize = c.msize

Loop:
	for c.Next() {
		for _, m := range c.Messages() {
			tver, ok := m.(styxproto.Tversion)
			if !ok {
				c.Rerror(m.Tag(), "need Tversion")
				break Loop
			}
			msize := tver.Msize()
			if msize < styxproto.MinBufSize {
				c.Rerror(m.Tag(), "buffer too small")
				break Loop
			}
			if msize < c.msize {
				c.msize = msize
				c.Encoder.MaxSize = msize
				c.Decoder.MaxSize = msize
			}
			if !bytes.HasPrefix(tver.Version(), []byte("9P2000")) {
				c.Rversion(uint32(c.msize), "unknown")
			}
			c.Rversion(uint32(c.msize), "9P2000")
			return true
		}
	}
	return false
}

// NOTE(droyo) consider a scenario where a malicious actor connects
// to the server that repeatedly spams Tauth requests. It can quickly
// use up resources on the server. Consider the following measures:
//
// - rate-limiting Tauth requests
// - Setting a per-connection session limit
// - close connections that have not established a session in N seconds
func (c *conn) handleTauth(cx context.Context, m styxproto.Tauth) bool {
	defer c.clearTag(m.Tag())
	if c.srv.Auth == nil {
		c.Rerror(m.Tag(), "%s", errNotSupported)
		return true
	}
	if _, ok := c.sessionFid.Get(m.Afid()); ok {
		c.Rerror(m.Tag(), "fid %x in use", m.Afid())
		return false
	}
	client, server := net.Pipe()
	ch := &Channel{
		Context:         c.cx,
		ReadWriteCloser: server,
	}
	rwc, err := styxfile.New(client)
	if err != nil {
		// This should never happen
		panic(err)
	}
	s := newSession(c, m)
	go func() {
		s.authC <- c.srv.Auth(ch, s.User, s.Access)
		close(s.authC)
	}()

	c.sessionFid.Put(m.Afid(), s)
	s.files.Put(m.Afid(), file{rwc: rwc, auth: true})
	s.IncRef()
	return true
}

func (c *conn) handleTattach(cx context.Context, m styxproto.Tattach) bool {
	defer c.clearTag(m.Tag())
	var handler Handler = DefaultServeMux
	if c.srv.Handler != nil {
		handler = c.srv.Handler
	}
	var s *Session
	if c.srv.Auth == nil {
		s = newSession(c, m)
	} else {
		// TODO(droyo) when a transport-based authentication scheme
		// is in use, the client should not have to do a Tauth request.
		// We should call the Auth handler if Afid is NOFID, passing it
		// a util.BlackHole.
		if !c.sessionFid.Fetch(s, m.Afid()) {
			c.Rerror(m.Tag(), "invalid afid %x", m.Afid())
			return false
		}
		// From attach(5): The same validated afid may be used for
		// multiple attach messages with the same uname and aname.
		if s.User != string(m.Uname()) || s.Access != string(m.Aname()) {
			c.Rerror(m.Tag(), "afid mismatch for %s on %s", m.Uname(), m.Aname())
			return false
		}
		if err := <-s.authC; err != nil {
			c.Rerror(m.Tag(), "auth failed: %s", err)
			return false
		}
	}
	go func() {
		handler.Serve9P(s)
		s.cleanupHandler()
	}()
	c.sessionFid.Put(m.Fid(), s)
	s.IncRef()
	s.files.Put(m.Fid(), file{name: "/", rwc: nil})
	c.Rattach(m.Tag(), c.qid(".", styxproto.QTDIR))
	return true
}

func (c *conn) handleTflush(cx context.Context, m styxproto.Tflush) bool {
	defer c.clearTag(m.Tag())

	oldtag := m.Oldtag()
	c.clearTag(oldtag)

	c.Rflush(m.Tag())
	return true
}

func (c *conn) handleFcall(cx context.Context, msg fcall) bool {
	s, ok := c.sessionByFid(msg.Fid())
	if !ok {
		c.Rerror(msg.Tag(), "%s", errNoFid)
		return false
	}

	file, ok := s.fetchFile(msg.Fid())
	if !ok {
		panic("bug: fid in session map, but no file associated")
		return false
	}

	// NOTE(droyo) on security and anonymous users: On a server with
	// authentication enabled, a client can only ever establish a handle
	// to the auth file.  At this point, we have checked that the fid
	// is valid, so *file can only be an auth file if the user has not
	// completed a Tattach.
	if file.auth {
		// Limit the number of request handlers we have to
		// audit.
		switch msg := msg.(type) {
		case styxproto.Twrite:
		case styxproto.Tread:
		case styxproto.Tstat:
		case styxproto.Tclunk:
		default:
			c.Rerror(msg.Tag(), "%T not allowed on afid", msg)
			c.clearTag(msg.Tag())
			return false
		}
	}

	switch msg := msg.(type) {
	case styxproto.Twalk:
		return s.handleTwalk(cx, msg, file)
	case styxproto.Topen:
		return s.handleTopen(cx, msg, file)
	case styxproto.Tcreate:
		return s.handleTcreate(cx, msg, file)
	case styxproto.Tread:
		return s.handleTread(cx, msg, file)
	case styxproto.Twrite:
		return s.handleTwrite(cx, msg, file)
	case styxproto.Tremove:
		return s.handleTremove(cx, msg, file)
	case styxproto.Tstat:
		return s.handleTstat(cx, msg, file)
	case styxproto.Twstat:
		return s.handleTwstat(cx, msg, file)
	case styxproto.Tclunk:
		return s.handleTclunk(cx, msg, file)
	}
	// invalid messages should have been caught
	// in the conn.serve loop, so we should never
	// reach this point.
	panic(fmt.Errorf("unhandled message type %T", msg))
}
<commit_msg>Protect pendingReq map from concurrent access<commit_after>package styx

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"

	"aqwari.net/net/styx/internal/qidpool"
	"aqwari.net/net/styx/internal/styxfile"
	"aqwari.net/net/styx/internal/tracing"
	"aqwari.net/net/styx/internal/util"
	"aqwari.net/net/styx/styxproto"

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

var (
	errFidInUse     = errors.New("fid already in use")
	errTagInUse     = errors.New("tag in use")
	errNoFid        = errors.New("no such fid")
	errNotSupported = errors.New("not supported")
)

type fcall interface {
	styxproto.Msg
	Fid() uint32
}

// A note on identifiers (fids & tags)
//
// identifiers are chosen by the client, not by the server.  Therefore,
// it is important that the performance and behavior of a server does
// *not* change based on the fid or tag a client chooses. This is why
// a map is used; its performance is good, and doesn't change based
// on the input a client chooses (rather, it does not change in a way
// a client can predict).

// A conn receives and sends 9P messages across a single network connection.
// Multiple "sessions" may take place over a single connection. The conn
// struct contains the necessary information to route 9P messages to their
// established sessions.
type conn struct {
	// These wrap the network connection to read and write messages.
	*styxproto.Decoder
	*styxproto.Encoder

	// The Server a connection was spawned from. Contains configuration
	// settings and the authentication function, if any.
	srv *Server

	// The network connection itself. We expose it in the struct so that
	// it is available for transport-based auth and any timeouts we need
	// to implement.
	rwc io.ReadWriteCloser

	// This serves as the parent context for the context attached to all
	// requests.
	cx context.Context

	// While srv.MaxSize holds the *desired* 9P protocol message
	// size, msize will contain the actual maximum negotiated with
	// the client, through a Tversion/Rversion exchange.
	msize int64

	// There is no "session id" in 9P. However, because all fids
	// for a connection must be derived from the fid established
	// in a Tattach call, any message that contains a fid can be
	// traced back to the original Tattach message.
	sessionFid *util.Map

	// Qids for the file tree, added on-demand.
	qidpool *qidpool.Pool

	// used to implement request cancellation when a Tflush
	// message is received.
	pendingReq *util.Map
}

func (c *conn) remoteAddr() net.Addr {
	type hasRemote interface {
		RemoteAddr() net.Addr
	}
	if nc, ok := c.rwc.(hasRemote); ok {
		return nc.RemoteAddr()
	}
	return nil
}

func (c *conn) sessionByFid(fid uint32) (*Session, bool) {
	if v, ok := c.sessionFid.Get(fid); ok {
		return v.(*Session), true
	}
	return nil, false
}

// Close the connection
func (c *conn) close() error {
	// Cancel all pending requests
	c.pendingReq.Do(func(m map[interface{}]interface{}) {
		for tag, cancel := range m {
			cancel.(context.CancelFunc)()
			delete(m, tag)
		}
	})

	// Close all open files and sessions
	c.sessionFid.Do(func(m map[interface{}]interface{}) {
		seen := make(map[*Session]struct{}, len(m))
		for k, v := range m {
			session := v.(*Session)
			if _, ok := seen[session]; !ok {
				seen[session] = struct{}{}
				session.endSession()
			}
			// Should probably let the GC take care of this
			delete(m, k)
		}
	})

	return c.rwc.Close()
}

func newConn(srv *Server, rwc io.ReadWriteCloser) *conn {
	var msize int64 = styxproto.DefaultMaxSize
	if srv.MaxSize > 0 {
		if srv.MaxSize > styxproto.MinBufSize {
			msize = srv.MaxSize
		} else {
			msize = styxproto.MinBufSize
		}
	}
	var enc *styxproto.Encoder
	var dec *styxproto.Decoder
	if srv.TraceLog != nil {
		enc = tracing.Encoder(rwc, func(m styxproto.Msg) {
			srv.TraceLog.Printf("← %03d %s", m.Tag(), m)
		})
		dec = tracing.Decoder(rwc, func(m styxproto.Msg) {
			srv.TraceLog.Printf("→ %03d %s", m.Tag(), m)
		})
	} else {
		enc = styxproto.NewEncoder(rwc)
		dec = styxproto.NewDecoder(rwc)
	}
	return &conn{
		Decoder:    dec,
		Encoder:    enc,
		srv:        srv,
		rwc:        rwc,
		cx:         context.TODO(),
		msize:      msize,
		sessionFid: util.NewMap(),
		pendingReq: util.NewMap(),
		qidpool:    qidpool.New(),
	}
}

func (c *conn) qid(name string, qtype uint8) styxproto.Qid {
	return c.qidpool.Put(name, qtype)
}

// All request contexts must have their cancel functions
// called, to free up resources in the context.
func (c *conn) clearTag(tag uint16) {
	var cancel context.CancelFunc
	if c.pendingReq.Fetch(tag, &cancel) {
		cancel()
		c.pendingReq.Del(tag)
	}
}

// runs in its own goroutine, one per connection.
func (c *conn) serve() {
	defer c.close()

	if !c.acceptTversion() {
		return
	}

Loop:
	for c.Next() {
		for _, m := range c.Messages() {
			if !c.handleMessage(m) {
				break Loop
			}
		}
	}
	c.srv.logf("closed connection from %s", c.remoteAddr())
}

func (c *conn) handleMessage(m styxproto.Msg) bool {
	if _, ok := c.pendingReq.Get(m.Tag()); ok {
		c.Rerror(m.Tag(), "%s", errTagInUse)
		return false
	}
	cx, cancel := context.WithCancel(c.cx)
	c.pendingReq.Put(m.Tag(), cancel)

	switch m := m.(type) {
	case styxproto.Tauth:
		return c.handleTauth(cx, m)
	case styxproto.Tattach:
		return c.handleTattach(cx, m)
	case styxproto.Tflush:
		return c.handleTflush(cx, m)
	case fcall:
		return c.handleFcall(cx, m)
	case styxproto.BadMessage:
		c.srv.logf("got bad message from %s: %s", c.remoteAddr(), m.Err)
		c.Rerror(m.Tag(), "bad message: %s", m.Err)
		c.clearTag(m.Tag())
		return true
	default:
		c.Rerror(m.Tag(), "unexpected %T message", m)
		return false
	}
	return true
}

// This is the first thing we do on a new connection. The first
// message a client sends *must* be a Tversion message.
func (c *conn) acceptTversion() bool {
	c.Encoder.MaxSize = c.msize
	c.Decoder.MaxSize = c.msize

Loop:
	for c.Next() {
		for _, m := range c.Messages() {
			tver, ok := m.(styxproto.Tversion)
			if !ok {
				c.Rerror(m.Tag(), "need Tversion")
				break Loop
			}
			msize := tver.Msize()
			if msize < styxproto.MinBufSize {
				c.Rerror(m.Tag(), "buffer too small")
				break Loop
			}
			if msize < c.msize {
				c.msize = msize
				c.Encoder.MaxSize = msize
				c.Decoder.MaxSize = msize
			}
			if !bytes.HasPrefix(tver.Version(), []byte("9P2000")) {
				c.Rversion(uint32(c.msize), "unknown")
			}
			c.Rversion(uint32(c.msize), "9P2000")
			return true
		}
	}
	return false
}

// NOTE(droyo) consider a scenario where a malicious actor connects
// to the server that repeatedly spams Tauth requests. It can quickly
// use up resources on the server. Consider the following measures:
//
// - rate-limiting Tauth requests
// - Setting a per-connection session limit
// - close connections that have not established a session in N seconds
func (c *conn) handleTauth(cx context.Context, m styxproto.Tauth) bool {
	defer c.clearTag(m.Tag())
	if c.srv.Auth == nil {
		c.Rerror(m.Tag(), "%s", errNotSupported)
		return true
	}
	if _, ok := c.sessionFid.Get(m.Afid()); ok {
		c.Rerror(m.Tag(), "fid %x in use", m.Afid())
		return false
	}
	client, server := net.Pipe()
	ch := &Channel{
		Context:         c.cx,
		ReadWriteCloser: server,
	}
	rwc, err := styxfile.New(client)
	if err != nil {
		// This should never happen
		panic(err)
	}
	s := newSession(c, m)
	go func() {
		s.authC <- c.srv.Auth(ch, s.User, s.Access)
		close(s.authC)
	}()

	c.sessionFid.Put(m.Afid(), s)
	s.files.Put(m.Afid(), file{rwc: rwc, auth: true})
	s.IncRef()
	return true
}

func (c *conn) handleTattach(cx context.Context, m styxproto.Tattach) bool {
	defer c.clearTag(m.Tag())
	var handler Handler = DefaultServeMux
	if c.srv.Handler != nil {
		handler = c.srv.Handler
	}
	var s *Session
	if c.srv.Auth == nil {
		s = newSession(c, m)
	} else {
		// TODO(droyo) when a transport-based authentication scheme
		// is in use, the client should not have to do a Tauth request.
		// We should call the Auth handler if Afid is NOFID, passing it
		// a util.BlackHole.
		if !c.sessionFid.Fetch(s, m.Afid()) {
			c.Rerror(m.Tag(), "invalid afid %x", m.Afid())
			return false
		}
		// From attach(5): The same validated afid may be used for
		// multiple attach messages with the same uname and aname.
		if s.User != string(m.Uname()) || s.Access != string(m.Aname()) {
			c.Rerror(m.Tag(), "afid mismatch for %s on %s", m.Uname(), m.Aname())
			return false
		}
		if err := <-s.authC; err != nil {
			c.Rerror(m.Tag(), "auth failed: %s", err)
			return false
		}
	}
	go func() {
		handler.Serve9P(s)
		s.cleanupHandler()
	}()
	c.sessionFid.Put(m.Fid(), s)
	s.IncRef()
	s.files.Put(m.Fid(), file{name: "/", rwc: nil})
	c.Rattach(m.Tag(), c.qid(".", styxproto.QTDIR))
	return true
}

func (c *conn) handleTflush(cx context.Context, m styxproto.Tflush) bool {
	defer c.clearTag(m.Tag())

	oldtag := m.Oldtag()
	c.clearTag(oldtag)

	c.Rflush(m.Tag())
	return true
}

func (c *conn) handleFcall(cx context.Context, msg fcall) bool {
	s, ok := c.sessionByFid(msg.Fid())
	if !ok {
		c.Rerror(msg.Tag(), "%s", errNoFid)
		return false
	}

	file, ok := s.fetchFile(msg.Fid())
	if !ok {
		panic("bug: fid in session map, but no file associated")
		return false
	}

	// NOTE(droyo) on security and anonymous users: On a server with
	// authentication enabled, a client can only ever establish a handle
	// to the auth file.  At this point, we have checked that the fid
	// is valid, so *file can only be an auth file if the user has not
	// completed a Tattach.
	if file.auth {
		// Limit the number of request handlers we have to
		// audit.
		switch msg := msg.(type) {
		case styxproto.Twrite:
		case styxproto.Tread:
		case styxproto.Tstat:
		case styxproto.Tclunk:
		default:
			c.Rerror(msg.Tag(), "%T not allowed on afid", msg)
			c.clearTag(msg.Tag())
			return false
		}
	}

	switch msg := msg.(type) {
	case styxproto.Twalk:
		return s.handleTwalk(cx, msg, file)
	case styxproto.Topen:
		return s.handleTopen(cx, msg, file)
	case styxproto.Tcreate:
		return s.handleTcreate(cx, msg, file)
	case styxproto.Tread:
		return s.handleTread(cx, msg, file)
	case styxproto.Twrite:
		return s.handleTwrite(cx, msg, file)
	case styxproto.Tremove:
		return s.handleTremove(cx, msg, file)
	case styxproto.Tstat:
		return s.handleTstat(cx, msg, file)
	case styxproto.Twstat:
		return s.handleTwstat(cx, msg, file)
	case styxproto.Tclunk:
		return s.handleTclunk(cx, msg, file)
	}
	// invalid messages should have been caught
	// in the conn.serve loop, so we should never
	// reach this point.
	panic(fmt.Errorf("unhandled message type %T", msg))
}
<|endoftext|>


<commit_before>// Package udp provides a connection-oriented listener over a UDP PacketConn
package udp

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/dtls/v2/internal/net/deadline"
	"github.com/pion/transport/packetio"
)

const receiveMTU = 8192
const defaultListenBacklog = 128 // same as Linux default

var errClosedListener = errors.New("udp: listener closed")
var errListenQueueExceeded = errors.New("udp: listen queue exceeded")

// listener augments a connection-oriented Listener over a UDP PacketConn
type listener struct {
	pConn *net.UDPConn

	accepting atomic.Value // bool
	acceptCh  chan *Conn
	doneCh    chan struct{}
	doneOnce  sync.Once

	connLock sync.Mutex
	conns    map[string]*Conn
	connWG   sync.WaitGroup

	readWG   sync.WaitGroup
	errClose atomic.Value // error
}

// Accept waits for and returns the next connection to the listener.
func (l *listener) Accept() (net.Conn, error) {
	select {
	case c := <-l.acceptCh:
		l.connWG.Add(1)
		return c, nil

	case <-l.doneCh:
		return nil, errClosedListener
	}
}

// Close closes the listener.
// Any blocked Accept operations will be unblocked and return errors.
func (l *listener) Close() error {
	var err error
	l.doneOnce.Do(func() {
		l.accepting.Store(false)
		close(l.doneCh)

		l.connLock.Lock()
		// Close unaccepted connections
	L_CLOSE:
		for {
			select {
			case c := <-l.acceptCh:
				close(c.doneCh)
				delete(l.conns, c.rAddr.String())

			default:
				break L_CLOSE
			}
		}
		nConns := len(l.conns)
		l.connLock.Unlock()

		l.connWG.Done()

		if nConns == 0 {
			// Wait if this is the final connection
			l.readWG.Wait()
			if errClose, ok := l.errClose.Load().(error); ok {
				err = errClose
			}
		} else {
			err = nil
		}
	})

	return err
}

// Addr returns the listener's network address.
func (l *listener) Addr() net.Addr {
	return l.pConn.LocalAddr()
}

// ListenConfig stores options for listening to an address.
type ListenConfig struct {
	// Backlog defines the maximum length of the queue of pending
	// connections. It is equivalent of the backlog argument of
	// POSIX listen function.
	// If a connection request arrives when the queue is full,
	// the request will be silently discarded, unlike TCP.
	// Set zero to use default value 128 which is same as Linux default.
	Backlog int
}

// Listen creates a new listener based on the ListenConfig.
func (lc *ListenConfig) Listen(network string, laddr *net.UDPAddr) (net.Listener, error) {
	if lc.Backlog == 0 {
		lc.Backlog = defaultListenBacklog
	}

	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}

	l := &listener{
		pConn:    conn,
		acceptCh: make(chan *Conn, lc.Backlog),
		conns:    make(map[string]*Conn),
		doneCh:   make(chan struct{}),
	}
	l.accepting.Store(true)
	l.connWG.Add(1)
	l.readWG.Add(2) // wait readLoop and Close execution routine

	go l.readLoop()
	go func() {
		l.connWG.Wait()
		if err := l.pConn.Close(); err != nil {
			l.errClose.Store(err)
		}
		l.readWG.Done()
	}()

	return l, nil
}

// Listen creates a new listener using default ListenConfig.
func Listen(network string, laddr *net.UDPAddr) (net.Listener, error) {
	return (&ListenConfig{}).Listen(network, laddr)
}

var readBufferPool = &sync.Pool{
	New: func() interface{} {
		buf := make([]byte, receiveMTU)
		return &buf
	},
}

// readLoop has to tasks:
// 1. Dispatching incoming packets to the correct Conn.
//    It can therefore not be ended until all Conns are closed.
// 2. Creating a new Conn when receiving from a new remote.
func (l *listener) readLoop() {
	defer l.readWG.Done()

	for {
		buf := *(readBufferPool.Get().(*[]byte))
		n, raddr, err := l.pConn.ReadFrom(buf)
		if err != nil {
			return
		}
		conn, err := l.getConn(raddr)
		if err != nil {
			continue
		}
		_, _ = conn.buffer.Write(buf[:n])
	}
}

func (l *listener) getConn(raddr net.Addr) (*Conn, error) {
	l.connLock.Lock()
	defer l.connLock.Unlock()
	conn, ok := l.conns[raddr.String()]
	if !ok {
		if !l.accepting.Load().(bool) {
			return nil, errClosedListener
		}
		conn = l.newConn(raddr)
		select {
		case l.acceptCh <- conn:
			l.conns[raddr.String()] = conn
		default:
			return nil, errListenQueueExceeded
		}
	}
	return conn, nil
}

// Conn augments a connection-oriented connection over a UDP PacketConn
type Conn struct {
	listener *listener

	rAddr net.Addr

	buffer *packetio.Buffer

	doneCh   chan struct{}
	doneOnce sync.Once

	writeDeadline *deadline.Deadline
}

func (l *listener) newConn(rAddr net.Addr) *Conn {
	return &Conn{
		listener:      l,
		rAddr:         rAddr,
		buffer:        packetio.NewBuffer(),
		doneCh:        make(chan struct{}),
		writeDeadline: deadline.New(),
	}
}

// Read
func (c *Conn) Read(p []byte) (int, error) {
	return c.buffer.Read(p)
}

// Write writes len(p) bytes from p to the DTLS connection
func (c *Conn) Write(p []byte) (n int, err error) {
	select {
	case <-c.writeDeadline.Done():
		return 0, context.DeadlineExceeded
	default:
	}
	return c.listener.pConn.WriteTo(p, c.rAddr)
}

// Close closes the conn and releases any Read calls
func (c *Conn) Close() error {
	var err error
	c.doneOnce.Do(func() {
		c.listener.connWG.Done()
		close(c.doneCh)
		c.listener.connLock.Lock()
		delete(c.listener.conns, c.rAddr.String())
		nConns := len(c.listener.conns)
		c.listener.connLock.Unlock()

		if nConns == 0 && !c.listener.accepting.Load().(bool) {
			// Wait if this is the final connection
			c.listener.readWG.Wait()
			if errClose, ok := c.listener.errClose.Load().(error); ok {
				err = errClose
			}
		} else {
			err = nil
		}
	})

	return err
}

// LocalAddr implements net.Conn.LocalAddr
func (c *Conn) LocalAddr() net.Addr {
	return c.listener.pConn.LocalAddr()
}

// RemoteAddr implements net.Conn.RemoteAddr
func (c *Conn) RemoteAddr() net.Addr {
	return c.rAddr
}

// SetDeadline implements net.Conn.SetDeadline
func (c *Conn) SetDeadline(t time.Time) error {
	c.writeDeadline.Set(t)
	return c.SetReadDeadline(t)
}

// SetReadDeadline implements net.Conn.SetDeadline
func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.buffer.SetReadDeadline(t)
}

// SetWriteDeadline implements net.Conn.SetDeadline
func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.writeDeadline.Set(t)
	// Write deadline of underlying connection should not be changed
	// since the connection can be shared.
	return nil
}
<commit_msg>Update module pion/transport to v0.10.0<commit_after>// Package udp provides a connection-oriented listener over a UDP PacketConn
package udp

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/transport/deadline"
	"github.com/pion/transport/packetio"
)

const receiveMTU = 8192
const defaultListenBacklog = 128 // same as Linux default

var errClosedListener = errors.New("udp: listener closed")
var errListenQueueExceeded = errors.New("udp: listen queue exceeded")

// listener augments a connection-oriented Listener over a UDP PacketConn
type listener struct {
	pConn *net.UDPConn

	accepting atomic.Value // bool
	acceptCh  chan *Conn
	doneCh    chan struct{}
	doneOnce  sync.Once

	connLock sync.Mutex
	conns    map[string]*Conn
	connWG   sync.WaitGroup

	readWG   sync.WaitGroup
	errClose atomic.Value // error
}

// Accept waits for and returns the next connection to the listener.
func (l *listener) Accept() (net.Conn, error) {
	select {
	case c := <-l.acceptCh:
		l.connWG.Add(1)
		return c, nil

	case <-l.doneCh:
		return nil, errClosedListener
	}
}

// Close closes the listener.
// Any blocked Accept operations will be unblocked and return errors.
func (l *listener) Close() error {
	var err error
	l.doneOnce.Do(func() {
		l.accepting.Store(false)
		close(l.doneCh)

		l.connLock.Lock()
		// Close unaccepted connections
	L_CLOSE:
		for {
			select {
			case c := <-l.acceptCh:
				close(c.doneCh)
				delete(l.conns, c.rAddr.String())

			default:
				break L_CLOSE
			}
		}
		nConns := len(l.conns)
		l.connLock.Unlock()

		l.connWG.Done()

		if nConns == 0 {
			// Wait if this is the final connection
			l.readWG.Wait()
			if errClose, ok := l.errClose.Load().(error); ok {
				err = errClose
			}
		} else {
			err = nil
		}
	})

	return err
}

// Addr returns the listener's network address.
func (l *listener) Addr() net.Addr {
	return l.pConn.LocalAddr()
}

// ListenConfig stores options for listening to an address.
type ListenConfig struct {
	// Backlog defines the maximum length of the queue of pending
	// connections. It is equivalent of the backlog argument of
	// POSIX listen function.
	// If a connection request arrives when the queue is full,
	// the request will be silently discarded, unlike TCP.
	// Set zero to use default value 128 which is same as Linux default.
	Backlog int
}

// Listen creates a new listener based on the ListenConfig.
func (lc *ListenConfig) Listen(network string, laddr *net.UDPAddr) (net.Listener, error) {
	if lc.Backlog == 0 {
		lc.Backlog = defaultListenBacklog
	}

	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}

	l := &listener{
		pConn:    conn,
		acceptCh: make(chan *Conn, lc.Backlog),
		conns:    make(map[string]*Conn),
		doneCh:   make(chan struct{}),
	}
	l.accepting.Store(true)
	l.connWG.Add(1)
	l.readWG.Add(2) // wait readLoop and Close execution routine

	go l.readLoop()
	go func() {
		l.connWG.Wait()
		if err := l.pConn.Close(); err != nil {
			l.errClose.Store(err)
		}
		l.readWG.Done()
	}()

	return l, nil
}

// Listen creates a new listener using default ListenConfig.
func Listen(network string, laddr *net.UDPAddr) (net.Listener, error) {
	return (&ListenConfig{}).Listen(network, laddr)
}

var readBufferPool = &sync.Pool{
	New: func() interface{} {
		buf := make([]byte, receiveMTU)
		return &buf
	},
}

// readLoop has to tasks:
// 1. Dispatching incoming packets to the correct Conn.
//    It can therefore not be ended until all Conns are closed.
// 2. Creating a new Conn when receiving from a new remote.
func (l *listener) readLoop() {
	defer l.readWG.Done()

	for {
		buf := *(readBufferPool.Get().(*[]byte))
		n, raddr, err := l.pConn.ReadFrom(buf)
		if err != nil {
			return
		}
		conn, err := l.getConn(raddr)
		if err != nil {
			continue
		}
		_, _ = conn.buffer.Write(buf[:n])
	}
}

func (l *listener) getConn(raddr net.Addr) (*Conn, error) {
	l.connLock.Lock()
	defer l.connLock.Unlock()
	conn, ok := l.conns[raddr.String()]
	if !ok {
		if !l.accepting.Load().(bool) {
			return nil, errClosedListener
		}
		conn = l.newConn(raddr)
		select {
		case l.acceptCh <- conn:
			l.conns[raddr.String()] = conn
		default:
			return nil, errListenQueueExceeded
		}
	}
	return conn, nil
}

// Conn augments a connection-oriented connection over a UDP PacketConn
type Conn struct {
	listener *listener

	rAddr net.Addr

	buffer *packetio.Buffer

	doneCh   chan struct{}
	doneOnce sync.Once

	writeDeadline *deadline.Deadline
}

func (l *listener) newConn(rAddr net.Addr) *Conn {
	return &Conn{
		listener:      l,
		rAddr:         rAddr,
		buffer:        packetio.NewBuffer(),
		doneCh:        make(chan struct{}),
		writeDeadline: deadline.New(),
	}
}

// Read
func (c *Conn) Read(p []byte) (int, error) {
	return c.buffer.Read(p)
}

// Write writes len(p) bytes from p to the DTLS connection
func (c *Conn) Write(p []byte) (n int, err error) {
	select {
	case <-c.writeDeadline.Done():
		return 0, context.DeadlineExceeded
	default:
	}
	return c.listener.pConn.WriteTo(p, c.rAddr)
}

// Close closes the conn and releases any Read calls
func (c *Conn) Close() error {
	var err error
	c.doneOnce.Do(func() {
		c.listener.connWG.Done()
		close(c.doneCh)
		c.listener.connLock.Lock()
		delete(c.listener.conns, c.rAddr.String())
		nConns := len(c.listener.conns)
		c.listener.connLock.Unlock()

		if nConns == 0 && !c.listener.accepting.Load().(bool) {
			// Wait if this is the final connection
			c.listener.readWG.Wait()
			if errClose, ok := c.listener.errClose.Load().(error); ok {
				err = errClose
			}
		} else {
			err = nil
		}
	})

	return err
}

// LocalAddr implements net.Conn.LocalAddr
func (c *Conn) LocalAddr() net.Addr {
	return c.listener.pConn.LocalAddr()
}

// RemoteAddr implements net.Conn.RemoteAddr
func (c *Conn) RemoteAddr() net.Addr {
	return c.rAddr
}

// SetDeadline implements net.Conn.SetDeadline
func (c *Conn) SetDeadline(t time.Time) error {
	c.writeDeadline.Set(t)
	return c.SetReadDeadline(t)
}

// SetReadDeadline implements net.Conn.SetDeadline
func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.buffer.SetReadDeadline(t)
}

// SetWriteDeadline implements net.Conn.SetDeadline
func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.writeDeadline.Set(t)
	// Write deadline of underlying connection should not be changed
	// since the connection can be shared.
	return nil
}
<|endoftext|>