<commit_before>package main

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
}

var iIndex *InvertedIndex
var fIndex *ForwardIndex

func main() {
	iIndex = NewInvertedIndex()
	fIndex = NewForwardIndex()

	InitIndex(iIndex, fIndex)

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

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("w1_fixed.txt")

	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
}

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
}

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
}

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)
)

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]

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

		hash ^= uint64(0xFF & (c >> 16))
		hash *= FNV_PRIME_64

		bitpos := hash % NUM_BITS
		if bitpos < 0 {
			bitpos += NUM_BITS
		}
		filter = filter | (1 << bitpos)
	}

	return filter
}

//Inverted Index
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
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 the port and word text a passed in paramter<commit_after>/*
 * Copyright (c) 2011 Jamra.source@gmail.com
 *
 * 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]
}
<|endoftext|>


<commit_before>package conn

import (
	"errors"
	"github.com/gorilla/websocket"
	"net/http"
	"sync"
	"time"
)

var (
	// ErrMessageChannelFull is returned when the connection's envelope channel is full.
	ErrMessageChannelFull = errors.New("Message channel is full")

	closeEnvelope          = &envelope{websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")}
	closeGoingAwayEnvelope = &envelope{websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "")}
	pingEnvelope           = &envelope{websocket.PingMessage, []byte{}}
	pongEnvelope           = &envelope{websocket.PongMessage, []byte{}}
)

// Conn represents a web socket connection.
type Conn struct {
	Settings             *Settings
	BinaryMessageHandler func([]byte)
	TextMessageHandler   func(string)
	DisconnectHandler    func()
	ErrorHandler         func(error)

	conn       *websocket.Conn
	envelopeCh chan *envelope
	mu         *sync.Mutex
}

// New creates a Conn.
func New() *Conn {
	return &Conn{
		Settings: newDefaultSettings(),
		mu:       new(sync.Mutex),
	}
}

// Connect to the peer.
func (c *Conn) Connect(url string, requestHeader http.Header) (*http.Response, error) {
	c.envelopeCh = make(chan *envelope, c.Settings.MessageChannelBufferSize)

	dialer := new(websocket.Dialer)
	dialer.ReadBufferSize = c.Settings.ReadBufferSize
	dialer.WriteBufferSize = c.Settings.WriteBufferSize
	dialer.HandshakeTimeout = c.Settings.HandshakeTimeout
	dialer.Subprotocols = c.Settings.Subprotocols
	dialer.NetDial = c.Settings.DialerSettings.NetDial
	dialer.TLSClientConfig = c.Settings.DialerSettings.TLSClientConfig

	conn, response, err := dialer.Dial(url, requestHeader)
	if err != nil {
		return response, err
	}
	c.conn = conn
	c.start()
	return response, nil
}

// UpgradeFromHTTP upgrades HTTP to WebSocket.
func (c *Conn) UpgradeFromHTTP(responseWriter http.ResponseWriter, request *http.Request) error {
	c.envelopeCh = make(chan *envelope, c.Settings.MessageChannelBufferSize)

	upgrader := new(websocket.Upgrader)
	upgrader.ReadBufferSize = c.Settings.ReadBufferSize
	upgrader.WriteBufferSize = c.Settings.WriteBufferSize
	upgrader.HandshakeTimeout = c.Settings.HandshakeTimeout
	upgrader.Subprotocols = c.Settings.Subprotocols
	upgrader.Error = c.Settings.UpgraderSettings.Error
	upgrader.CheckOrigin = c.Settings.UpgraderSettings.CheckOrigin

	conn, err := upgrader.Upgrade(responseWriter, request, nil)
	if err != nil {
		return err
	}
	c.conn = conn
	c.start()
	return nil
}

// Close the connection.
func (c *Conn) Close() error {
	return c.postEnvelope(closeEnvelope)
}

// WriteBinaryMessage to the peer.
func (c *Conn) WriteBinaryMessage(data []byte) error {
	return c.postEnvelope(&envelope{websocket.BinaryMessage, data})
}

// WriteTextMessage to the peer.
func (c *Conn) WriteTextMessage(text string) error {
	return c.postEnvelope(&envelope{websocket.TextMessage, []byte(text)})
}

func (c *Conn) start() {
	c.conn.SetReadLimit(c.Settings.MaxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(c.Settings.PongWait))
	c.conn.SetPingHandler(func(string) error {
		return c.postEnvelope(pongEnvelope)
	})
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(c.Settings.PongWait))
	})

	go c.writePump()
	go c.readPump()
}

func (c *Conn) postEnvelope(e *envelope) error {
	select {
	case c.envelopeCh <- e:
		return nil
	default:
		return ErrMessageChannelFull
	}
}

func (c *Conn) writeMessage(e *envelope) error {
	if err := c.conn.SetWriteDeadline(time.Now().Add(c.Settings.WriteWait)); err != nil {
		c.handleError(err)
		return err
	}
	if err := c.conn.WriteMessage(e.messageType, e.data); err != nil {
		c.handleError(err)
		return err
	}
	return nil
}

func (c *Conn) handleError(err error) {
	if err != nil && c.ErrorHandler != nil {
		c.mu.Lock()
		defer c.mu.Unlock()
		c.ErrorHandler(err)
	}
}

func (c *Conn) writePump() {
	defer c.conn.Close()

	ticker := time.NewTicker(c.Settings.PingPeriod)
	defer ticker.Stop()

loop:
	for {
		select {
		case e, ok := <-c.envelopeCh:
			if !ok {
				c.writeMessage(closeGoingAwayEnvelope)
				break loop
			}
			if err := c.writeMessage(e); err != nil {
				break loop
			}
		case <-ticker.C:
			if err := c.writeMessage(pingEnvelope); err != nil {
				break loop
			}
		}
	}
}

func (c *Conn) readPump() {
	defer c.conn.Close()

	for {
		messageType, data, err := c.conn.ReadMessage()
		if err != nil {
			c.handleError(err)
			break
		}

		switch messageType {
		case websocket.BinaryMessage:
			if c.BinaryMessageHandler != nil {
				c.BinaryMessageHandler(data)
			}
		case websocket.TextMessage:
			if c.TextMessageHandler != nil {
				c.TextMessageHandler(string(data))
			}
		}
	}

	if c.DisconnectHandler != nil {
		c.DisconnectHandler()
	}
}
<commit_msg>Add ErrAlreadyConnected<commit_after>package conn

import (
	"errors"
	"github.com/gorilla/websocket"
	"net/http"
	"sync"
	"time"
)

var (
	// ErrMessageChannelFull is returned when the connection's envelope channel is full.
	ErrMessageChannelFull = errors.New("Message channel is full")

	// ErrAlreadyConnected is returned when the connection is already connected.
	ErrAlreadyConnected = errors.New("Already connected")

	closeEnvelope          = &envelope{websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")}
	closeGoingAwayEnvelope = &envelope{websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "")}
	pingEnvelope           = &envelope{websocket.PingMessage, []byte{}}
	pongEnvelope           = &envelope{websocket.PongMessage, []byte{}}
)

// Conn represents a web socket connection.
type Conn struct {
	Settings             *Settings
	BinaryMessageHandler func([]byte)
	TextMessageHandler   func(string)
	DisconnectHandler    func()
	ErrorHandler         func(error)

	conn       *websocket.Conn
	envelopeCh chan *envelope
	mu         *sync.Mutex
}

// New creates a Conn.
func New() *Conn {
	return &Conn{
		Settings: newDefaultSettings(),
		mu:       new(sync.Mutex),
	}
}

// Connect to the peer.
func (c *Conn) Connect(url string, requestHeader http.Header) (*http.Response, error) {
	if c.conn != nil {
		return nil, ErrAlreadyConnected
	}
	dialer := new(websocket.Dialer)
	dialer.ReadBufferSize = c.Settings.ReadBufferSize
	dialer.WriteBufferSize = c.Settings.WriteBufferSize
	dialer.HandshakeTimeout = c.Settings.HandshakeTimeout
	dialer.Subprotocols = c.Settings.Subprotocols
	dialer.NetDial = c.Settings.DialerSettings.NetDial
	dialer.TLSClientConfig = c.Settings.DialerSettings.TLSClientConfig

	conn, response, err := dialer.Dial(url, requestHeader)
	if err != nil {
		return response, err
	}
	c.conn = conn
	c.start()
	return response, nil
}

// UpgradeFromHTTP upgrades HTTP to WebSocket.
func (c *Conn) UpgradeFromHTTP(responseWriter http.ResponseWriter, request *http.Request) error {
	if c.conn != nil {
		return ErrAlreadyConnected
	}
	upgrader := new(websocket.Upgrader)
	upgrader.ReadBufferSize = c.Settings.ReadBufferSize
	upgrader.WriteBufferSize = c.Settings.WriteBufferSize
	upgrader.HandshakeTimeout = c.Settings.HandshakeTimeout
	upgrader.Subprotocols = c.Settings.Subprotocols
	upgrader.Error = c.Settings.UpgraderSettings.Error
	upgrader.CheckOrigin = c.Settings.UpgraderSettings.CheckOrigin

	conn, err := upgrader.Upgrade(responseWriter, request, nil)
	if err != nil {
		return err
	}
	c.conn = conn
	c.start()
	return nil
}

// Close the connection.
func (c *Conn) Close() error {
	return c.postEnvelope(closeEnvelope)
}

// WriteBinaryMessage to the peer.
func (c *Conn) WriteBinaryMessage(data []byte) error {
	return c.postEnvelope(&envelope{websocket.BinaryMessage, data})
}

// WriteTextMessage to the peer.
func (c *Conn) WriteTextMessage(text string) error {
	return c.postEnvelope(&envelope{websocket.TextMessage, []byte(text)})
}

func (c *Conn) start() {
	c.envelopeCh = make(chan *envelope, c.Settings.MessageChannelBufferSize)

	c.conn.SetReadLimit(c.Settings.MaxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(c.Settings.PongWait))
	c.conn.SetPingHandler(func(string) error {
		return c.postEnvelope(pongEnvelope)
	})
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(c.Settings.PongWait))
	})

	go c.writePump()
	go c.readPump()
}

func (c *Conn) postEnvelope(e *envelope) error {
	select {
	case c.envelopeCh <- e:
		return nil
	default:
		return ErrMessageChannelFull
	}
}

func (c *Conn) writeMessage(e *envelope) error {
	if err := c.conn.SetWriteDeadline(time.Now().Add(c.Settings.WriteWait)); err != nil {
		c.handleError(err)
		return err
	}
	if err := c.conn.WriteMessage(e.messageType, e.data); err != nil {
		c.handleError(err)
		return err
	}
	return nil
}

func (c *Conn) handleError(err error) {
	if err != nil && c.ErrorHandler != nil {
		c.mu.Lock()
		defer c.mu.Unlock()
		c.ErrorHandler(err)
	}
}

func (c *Conn) writePump() {
	defer c.conn.Close()

	ticker := time.NewTicker(c.Settings.PingPeriod)
	defer ticker.Stop()

loop:
	for {
		select {
		case e, ok := <-c.envelopeCh:
			if !ok {
				c.writeMessage(closeGoingAwayEnvelope)
				break loop
			}
			if err := c.writeMessage(e); err != nil {
				break loop
			}
		case <-ticker.C:
			if err := c.writeMessage(pingEnvelope); err != nil {
				break loop
			}
		}
	}
}

func (c *Conn) readPump() {
	defer c.conn.Close()

	for {
		messageType, data, err := c.conn.ReadMessage()
		if err != nil {
			c.handleError(err)
			break
		}

		switch messageType {
		case websocket.BinaryMessage:
			if c.BinaryMessageHandler != nil {
				c.BinaryMessageHandler(data)
			}
		case websocket.TextMessage:
			if c.TextMessageHandler != nil {
				c.TextMessageHandler(string(data))
			}
		}
	}

	if c.DisconnectHandler != nil {
		c.DisconnectHandler()
	}
}
<|endoftext|>


<commit_before>package smtp

import (
	"bufio"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Conn struct {
	server    *Server
	helo      string
	User      User
	msg       *Message
	conn      net.Conn
	reader    *bufio.Reader
	writer    *bufio.Writer
	nbrErrors int
}

func newConn(c net.Conn, s *Server) *Conn {
	sc := &Conn{
		server: s,
		conn:   c,
	}

	sc.init()
	return sc
}

func (c *Conn) init() {
	r := io.Reader(c.conn)
	w := io.Writer(c.conn)

	if c.server.Debug != nil {
		r = io.TeeReader(r, c.server.Debug)
		w = io.MultiWriter(w, c.server.Debug)
	}

	c.reader = bufio.NewReader(r)
	c.writer = bufio.NewWriter(w)
}

// Commands are dispatched to the appropriate handler functions.
func (c *Conn) handle(cmd string, arg string) {
	if cmd == "" {
		c.Write("500", "Speak up")
		return
	}

	switch cmd {
	case "SEND", "SOML", "SAML", "EXPN", "HELP", "TURN":
		// These commands are not implemented in any state
		c.Write("502", fmt.Sprintf("%v command not implemented", cmd))
	case "HELO", "EHLO":
		c.handleGreet((cmd == "EHLO"), arg)
	case "MAIL":
		c.handleMail(arg)
	case "RCPT":
		c.handleRcpt(arg)
	case "VRFY":
		c.Write("252", "Cannot VRFY user, but will accept message")
	case "NOOP":
		c.Write("250", "I have sucessfully done nothing")
	case "RSET": // Reset session
		c.reset()
		c.Write("250", "Session reset")
	case "DATA":
		c.handleData(arg)
	case "QUIT":
		c.Write("221", "Goodnight and good luck")
		c.Close()
	case "AUTH":
		c.handleAuth(arg)
	case "STARTTLS":
		c.handleStartTLS()
	default:
		c.Write("500", fmt.Sprintf("Syntax error, %v command unrecognized", cmd))

		c.nbrErrors++
		if c.nbrErrors > 3 {
			c.Write("500", "Too many unrecognized commands")
			c.Close()
		}
	}
}

func (c *Conn) Close() error {
	if c.User != nil {
		c.User.Logout()
	}

	return c.conn.Close()
}

// Check if this connection is encrypted.
func (c *Conn) IsTLS() bool {
	_, ok := c.conn.(*tls.Conn)
	return ok
}

// GREET state -> waiting for HELO
func (c *Conn) handleGreet(enhanced bool, arg string) {
	if !enhanced {
		domain, err := parseHelloArgument(arg)
		if err != nil {
			c.Write("501", "Domain/address argument required for HELO")
			return
		}
		c.helo = domain

		c.Write("250", fmt.Sprintf("Hello %s", domain))
	} else {
		domain, err := parseHelloArgument(arg)
		if err != nil {
			c.Write("501", "Domain/address argument required for EHLO")
			return
		}

		c.helo = domain

		caps := []string{}
		caps = append(caps, c.server.caps...)
		if c.server.TLSConfig != nil && !c.IsTLS() {
			caps = append(caps, "STARTTLS")
		}
		if c.IsTLS() || c.server.AllowInsecureAuth {
			authCap := "AUTH"
			for name, _ := range c.server.auths {
				authCap += " " + name
			}

			caps = append(caps, authCap)
		}
		if c.server.MaxMessageBytes > 0 {
			caps = append(caps, fmt.Sprintf("SIZE %v", c.server.MaxMessageBytes))
		}

		args := []string{"Hello " + domain}
		args = append(args, caps...)
		c.Write("250", args...)
	}
}

// READY state -> waiting for MAIL
func (c *Conn) handleMail(arg string) {
	if c.helo == "" {
		c.Write("502", "Please introduce yourself first.")
		return
	}
	if c.msg == nil {
		c.Write("502", "Please authenticate first.")
		return
	}

	// Match FROM, while accepting '>' as quoted pair and in double quoted strings
	// (?i) makes the regex case insensitive, (?:) is non-grouping sub-match
	re := regexp.MustCompile("(?i)^FROM:\\s*<((?:\\\\>|[^>])+|\"[^\"]+\"@[^>]+)>( [\\w= ]+)?$")
	m := re.FindStringSubmatch(arg)
	if m == nil {
		c.Write("501", "Was expecting MAIL arg syntax of FROM:<address>")
		return
	}

	from := m[1]

	// This is where the Conn may put BODY=8BITMIME, but we already
	// read the DATA as bytes, so it does not effect our processing.
	if m[2] != "" {
		args, err := parseArgs(m[2])
		if err != nil {
			c.Write("501", "Unable to parse MAIL ESMTP parameters")
			return
		}

		if args["SIZE"] != "" {
			size, err := strconv.ParseInt(args["SIZE"], 10, 32)
			if err != nil {
				c.Write("501", "Unable to parse SIZE as an integer")
				return
			}

			if c.server.MaxMessageBytes > 0 && int(size) > c.server.MaxMessageBytes {
				c.Write("552", "Max message size exceeded")
				return
			}
		}
	}

	c.msg.From = from
	c.Write("250", fmt.Sprintf("Roger, accepting mail from <%v>", from))
}

// MAIL state -> waiting for RCPTs followed by DATA
func (c *Conn) handleRcpt(arg string) {
	if c.msg == nil || c.msg.From == "" {
		c.Write("502", "Missing MAIL FROM command.")
		return
	}

	if (len(arg) < 4) || (strings.ToUpper(arg[0:3]) != "TO:") {
		c.Write("501", "Was expecting RCPT arg syntax of TO:<address>")
		return
	}

	// TODO: This trim is probably too forgiving
	recipient := strings.Trim(arg[3:], "<> ")

	if c.server.MaxRecipients > 0 && len(c.msg.To) >= c.server.MaxRecipients {
		c.Write("552", fmt.Sprintf("Maximum limit of %v recipients reached", c.server.MaxRecipients))
		return
	}

	c.msg.To = append(c.msg.To, recipient)
	c.Write("250", fmt.Sprintf("I'll make sure <%v> gets this", recipient))
}

func (c *Conn) handleAuth(arg string) {
	if c.helo == "" {
		c.Write("502", "Please introduce yourself first.")
		return
	}

	if arg == "" {
		c.Write("502", "Missing parameter")
		return
	}

	parts := strings.Fields(arg)
	mechanism := strings.ToUpper(parts[0])

	// Parse client initial response if there is one
	var ir []byte
	if len(parts) > 1 {
		var err error
		ir, err = base64.StdEncoding.DecodeString(parts[1])
		if err != nil {
			return
		}
	}

	newSasl, ok := c.server.auths[mechanism]
	if !ok {
		c.Write("504", "Unsupported authentication mechanism")
		return
	}

	sasl := newSasl(c)
	scanner := bufio.NewScanner(c.reader)

	response := ir
	for {
		challenge, done, err := sasl.Next(response)
		if err != nil {
			c.Write("454", err.Error())
			return
		}

		if done {
			break
		}

		encoded := ""
		if len(challenge) > 0 {
			encoded = base64.StdEncoding.EncodeToString(challenge)
		}
		c.Write("334", encoded)

		if !scanner.Scan() {
			return
		}

		encoded = scanner.Text()
		if encoded != "" {
			response, err = base64.StdEncoding.DecodeString(encoded)
			if err != nil {
				c.Write("454", "Invalid base64 data")
				return
			}
		}
	}

	if c.User != nil {
		c.Write("235", "Authentication succeeded")

		c.msg = &Message{}
	}
}

func (c *Conn) handleStartTLS() {
	if c.IsTLS() {
		c.Write("502", "Already running in TLS")
		return
	}

	if c.server.TLSConfig == nil {
		c.Write("502", "TLS not supported")
		return
	}

	c.Write("220", "Ready to start TLS")

	// Upgrade to TLS
	var tlsConn *tls.Conn
	tlsConn = tls.Server(c.conn, c.server.TLSConfig)

	if err := tlsConn.Handshake(); err != nil {
		c.Write("550", "Handshake error")
	}

	c.conn = tlsConn
	c.init()

	// Reset envelope as a new EHLO/HELO is required after STARTTLS
	c.reset()
}

// DATA
func (c *Conn) handleData(arg string) {
	if arg != "" {
		c.Write("501", "DATA command should not have any arguments")
		return
	}

	if c.msg == nil || c.msg.From == "" || len(c.msg.To) == 0 {
		c.Write("502", "Missing RCPT TO command.")
		return
	}

	// We have recipients, go to accept data
	c.Write("354", "Go ahead. End your data with <CR><LF>.<CR><LF>")

	c.msg.Data = newDataReader(c)
	if err := c.User.Send(c.msg); err != nil {
		if err, ok := err.(*smtpError); ok {
			c.Write(err.Code, err.Message)
		} else {
			c.Write("554", "Error: transaction failed, blame it on the weather: "+err.Error())
		}
	} else {
		c.Write("250", "Ok: queued")
	}

	c.reset()
}

func (c *Conn) Reject() {
	c.Write("421", "Too busy. Try again later.")
	c.Close()
}

func (c *Conn) greet() {
	c.Write("220", fmt.Sprintf("%v ESMTP Service Ready", c.server.Domain))
}

// Calculate the next read or write deadline based on MaxIdleSeconds.
func (c *Conn) nextDeadline() time.Time {
	if c.server.MaxIdleSeconds == 0 {
		return time.Time{} // No deadline
	}

	return time.Now().Add(time.Duration(c.server.MaxIdleSeconds) * time.Second)
}

func (c *Conn) Write(code string, text ...string) {
	// TODO: error handling

	c.conn.SetDeadline(c.nextDeadline())

	for i := 0; i < len(text)-1; i++ {
		c.writer.Write([]byte(code + "-" + text[i] + "\r\n"))
	}
	c.writer.Write([]byte(code + " " + text[len(text)-1] + "\r\n"))
	c.writer.Flush()
}

// Reads a line of input
func (c *Conn) readLine() (line string, err error) {
	if err = c.conn.SetReadDeadline(c.nextDeadline()); err != nil {
		return "", err
	}

	line, err = c.reader.ReadString('\n')
	if err != nil {
		return "", err
	}

	return line, nil
}

func (c *Conn) reset() {
	if c.User != nil {
		c.User.Logout()
	}

	c.helo = ""
	c.User = nil
	c.msg = nil
}
<commit_msg>add server function<commit_after>package smtp

import (
	"bufio"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Conn struct {
	server    *Server
	helo      string
	User      User
	msg       *Message
	conn      net.Conn
	reader    *bufio.Reader
	writer    *bufio.Writer
	nbrErrors int
}

func newConn(c net.Conn, s *Server) *Conn {
	sc := &Conn{
		server: s,
		conn:   c,
	}

	sc.init()
	return sc
}

func (c *Conn) init() {
	r := io.Reader(c.conn)
	w := io.Writer(c.conn)

	if c.server.Debug != nil {
		r = io.TeeReader(r, c.server.Debug)
		w = io.MultiWriter(w, c.server.Debug)
	}

	c.reader = bufio.NewReader(r)
	c.writer = bufio.NewWriter(w)
}

// Commands are dispatched to the appropriate handler functions.
func (c *Conn) handle(cmd string, arg string) {
	if cmd == "" {
		c.Write("500", "Speak up")
		return
	}

	switch cmd {
	case "SEND", "SOML", "SAML", "EXPN", "HELP", "TURN":
		// These commands are not implemented in any state
		c.Write("502", fmt.Sprintf("%v command not implemented", cmd))
	case "HELO", "EHLO":
		c.handleGreet((cmd == "EHLO"), arg)
	case "MAIL":
		c.handleMail(arg)
	case "RCPT":
		c.handleRcpt(arg)
	case "VRFY":
		c.Write("252", "Cannot VRFY user, but will accept message")
	case "NOOP":
		c.Write("250", "I have sucessfully done nothing")
	case "RSET": // Reset session
		c.reset()
		c.Write("250", "Session reset")
	case "DATA":
		c.handleData(arg)
	case "QUIT":
		c.Write("221", "Goodnight and good luck")
		c.Close()
	case "AUTH":
		c.handleAuth(arg)
	case "STARTTLS":
		c.handleStartTLS()
	default:
		c.Write("500", fmt.Sprintf("Syntax error, %v command unrecognized", cmd))

		c.nbrErrors++
		if c.nbrErrors > 3 {
			c.Write("500", "Too many unrecognized commands")
			c.Close()
		}
	}
}

func (c *Conn) Server() *Server {
	return c.server
}

func (c *Conn) Close() error {
	if c.User != nil {
		c.User.Logout()
	}

	return c.conn.Close()
}

// Check if this connection is encrypted.
func (c *Conn) IsTLS() bool {
	_, ok := c.conn.(*tls.Conn)
	return ok
}

// GREET state -> waiting for HELO
func (c *Conn) handleGreet(enhanced bool, arg string) {
	if !enhanced {
		domain, err := parseHelloArgument(arg)
		if err != nil {
			c.Write("501", "Domain/address argument required for HELO")
			return
		}
		c.helo = domain

		c.Write("250", fmt.Sprintf("Hello %s", domain))
	} else {
		domain, err := parseHelloArgument(arg)
		if err != nil {
			c.Write("501", "Domain/address argument required for EHLO")
			return
		}

		c.helo = domain

		caps := []string{}
		caps = append(caps, c.server.caps...)
		if c.server.TLSConfig != nil && !c.IsTLS() {
			caps = append(caps, "STARTTLS")
		}
		if c.IsTLS() || c.server.AllowInsecureAuth {
			authCap := "AUTH"
			for name, _ := range c.server.auths {
				authCap += " " + name
			}

			caps = append(caps, authCap)
		}
		if c.server.MaxMessageBytes > 0 {
			caps = append(caps, fmt.Sprintf("SIZE %v", c.server.MaxMessageBytes))
		}

		args := []string{"Hello " + domain}
		args = append(args, caps...)
		c.Write("250", args...)
	}
}

// READY state -> waiting for MAIL
func (c *Conn) handleMail(arg string) {
	if c.helo == "" {
		c.Write("502", "Please introduce yourself first.")
		return
	}
	if c.msg == nil {
		c.Write("502", "Please authenticate first.")
		return
	}

	// Match FROM, while accepting '>' as quoted pair and in double quoted strings
	// (?i) makes the regex case insensitive, (?:) is non-grouping sub-match
	re := regexp.MustCompile("(?i)^FROM:\\s*<((?:\\\\>|[^>])+|\"[^\"]+\"@[^>]+)>( [\\w= ]+)?$")
	m := re.FindStringSubmatch(arg)
	if m == nil {
		c.Write("501", "Was expecting MAIL arg syntax of FROM:<address>")
		return
	}

	from := m[1]

	// This is where the Conn may put BODY=8BITMIME, but we already
	// read the DATA as bytes, so it does not effect our processing.
	if m[2] != "" {
		args, err := parseArgs(m[2])
		if err != nil {
			c.Write("501", "Unable to parse MAIL ESMTP parameters")
			return
		}

		if args["SIZE"] != "" {
			size, err := strconv.ParseInt(args["SIZE"], 10, 32)
			if err != nil {
				c.Write("501", "Unable to parse SIZE as an integer")
				return
			}

			if c.server.MaxMessageBytes > 0 && int(size) > c.server.MaxMessageBytes {
				c.Write("552", "Max message size exceeded")
				return
			}
		}
	}

	c.msg.From = from
	c.Write("250", fmt.Sprintf("Roger, accepting mail from <%v>", from))
}

// MAIL state -> waiting for RCPTs followed by DATA
func (c *Conn) handleRcpt(arg string) {
	if c.msg == nil || c.msg.From == "" {
		c.Write("502", "Missing MAIL FROM command.")
		return
	}

	if (len(arg) < 4) || (strings.ToUpper(arg[0:3]) != "TO:") {
		c.Write("501", "Was expecting RCPT arg syntax of TO:<address>")
		return
	}

	// TODO: This trim is probably too forgiving
	recipient := strings.Trim(arg[3:], "<> ")

	if c.server.MaxRecipients > 0 && len(c.msg.To) >= c.server.MaxRecipients {
		c.Write("552", fmt.Sprintf("Maximum limit of %v recipients reached", c.server.MaxRecipients))
		return
	}

	c.msg.To = append(c.msg.To, recipient)
	c.Write("250", fmt.Sprintf("I'll make sure <%v> gets this", recipient))
}

func (c *Conn) handleAuth(arg string) {
	if c.helo == "" {
		c.Write("502", "Please introduce yourself first.")
		return
	}

	if arg == "" {
		c.Write("502", "Missing parameter")
		return
	}

	parts := strings.Fields(arg)
	mechanism := strings.ToUpper(parts[0])

	// Parse client initial response if there is one
	var ir []byte
	if len(parts) > 1 {
		var err error
		ir, err = base64.StdEncoding.DecodeString(parts[1])
		if err != nil {
			return
		}
	}

	newSasl, ok := c.server.auths[mechanism]
	if !ok {
		c.Write("504", "Unsupported authentication mechanism")
		return
	}

	sasl := newSasl(c)
	scanner := bufio.NewScanner(c.reader)

	response := ir
	for {
		challenge, done, err := sasl.Next(response)
		if err != nil {
			c.Write("454", err.Error())
			return
		}

		if done {
			break
		}

		encoded := ""
		if len(challenge) > 0 {
			encoded = base64.StdEncoding.EncodeToString(challenge)
		}
		c.Write("334", encoded)

		if !scanner.Scan() {
			return
		}

		encoded = scanner.Text()
		if encoded != "" {
			response, err = base64.StdEncoding.DecodeString(encoded)
			if err != nil {
				c.Write("454", "Invalid base64 data")
				return
			}
		}
	}

	if c.User != nil {
		c.Write("235", "Authentication succeeded")

		c.msg = &Message{}
	}
}

func (c *Conn) handleStartTLS() {
	if c.IsTLS() {
		c.Write("502", "Already running in TLS")
		return
	}

	if c.server.TLSConfig == nil {
		c.Write("502", "TLS not supported")
		return
	}

	c.Write("220", "Ready to start TLS")

	// Upgrade to TLS
	var tlsConn *tls.Conn
	tlsConn = tls.Server(c.conn, c.server.TLSConfig)

	if err := tlsConn.Handshake(); err != nil {
		c.Write("550", "Handshake error")
	}

	c.conn = tlsConn
	c.init()

	// Reset envelope as a new EHLO/HELO is required after STARTTLS
	c.reset()
}

// DATA
func (c *Conn) handleData(arg string) {
	if arg != "" {
		c.Write("501", "DATA command should not have any arguments")
		return
	}

	if c.msg == nil || c.msg.From == "" || len(c.msg.To) == 0 {
		c.Write("502", "Missing RCPT TO command.")
		return
	}

	// We have recipients, go to accept data
	c.Write("354", "Go ahead. End your data with <CR><LF>.<CR><LF>")

	c.msg.Data = newDataReader(c)
	if err := c.User.Send(c.msg); err != nil {
		if err, ok := err.(*smtpError); ok {
			c.Write(err.Code, err.Message)
		} else {
			c.Write("554", "Error: transaction failed, blame it on the weather: "+err.Error())
		}
	} else {
		c.Write("250", "Ok: queued")
	}

	c.reset()
}

func (c *Conn) Reject() {
	c.Write("421", "Too busy. Try again later.")
	c.Close()
}

func (c *Conn) greet() {
	c.Write("220", fmt.Sprintf("%v ESMTP Service Ready", c.server.Domain))
}

// Calculate the next read or write deadline based on MaxIdleSeconds.
func (c *Conn) nextDeadline() time.Time {
	if c.server.MaxIdleSeconds == 0 {
		return time.Time{} // No deadline
	}

	return time.Now().Add(time.Duration(c.server.MaxIdleSeconds) * time.Second)
}

func (c *Conn) Write(code string, text ...string) {
	// TODO: error handling

	c.conn.SetDeadline(c.nextDeadline())

	for i := 0; i < len(text)-1; i++ {
		c.writer.Write([]byte(code + "-" + text[i] + "\r\n"))
	}
	c.writer.Write([]byte(code + " " + text[len(text)-1] + "\r\n"))
	c.writer.Flush()
}

// Reads a line of input
func (c *Conn) readLine() (line string, err error) {
	if err = c.conn.SetReadDeadline(c.nextDeadline()); err != nil {
		return "", err
	}

	line, err = c.reader.ReadString('\n')
	if err != nil {
		return "", err
	}

	return line, nil
}

func (c *Conn) reset() {
	if c.User != nil {
		c.User.Logout()
	}

	c.helo = ""
	c.User = nil
	c.msg = nil
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	defaultWelcomeMessage = "Welcome to the Go FTP Server"
)

type Conn struct {
	conn          net.Conn
	controlReader *bufio.Reader
	controlWriter *bufio.Writer
	dataConn      DataSocket
	driver        Driver
	auth          Auth
	logger        *Logger
	server        *Server
	sessionId     string
	namePrefix    string
	reqUser       string
	user          string
	renameFrom    string
	lastFilePos   int64
	appendData    bool
}

func (conn *Conn) LoginUser() string {
	return conn.user
}

func (conn *Conn) IsLogin() bool {
	return len(conn.user) > 0
}

// returns a random 20 char string that can be used as a unique session ID
func newSessionId() string {
	hash := sha256.New()
	_, err := io.CopyN(hash, rand.Reader, 50)
	if err != nil {
		return "????????????????????"
	}
	md := hash.Sum(nil)
	mdStr := hex.EncodeToString(md)
	return mdStr[0:20]
}

// Serve starts an endless loop that reads FTP commands from the client and
// responds appropriately. terminated is a channel that will receive a true
// message when the connection closes. This loop will be running inside a
// goroutine, so use this channel to be notified when the connection can be
// cleaned up.
func (conn *Conn) Serve() {
	conn.logger.Print("Connection Established")
	// send welcome
	conn.writeMessage(220, conn.server.WelcomeMessage)
	// read commands
	for {
		line, err := conn.controlReader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				continue
			}

			conn.logger.Print(fmt.Sprintln("read error:", err))
			break
		}
		conn.receiveLine(line)
	}
	conn.Close()
	conn.logger.Print("Connection Terminated")
}

// Close will manually close this connection, even if the client isn't ready.
func (conn *Conn) Close() {
	conn.conn.Close()
	if conn.dataConn != nil {
		conn.dataConn.Close()
		conn.dataConn = nil
	}
}

// receiveLine accepts a single line FTP command and co-ordinates an
// appropriate response.
func (Conn *Conn) receiveLine(line string) {
	command, param := Conn.parseLine(line)
	Conn.logger.PrintCommand(command, param)
	cmdObj := commands[strings.ToUpper(command)]
	if cmdObj == nil {
		Conn.writeMessage(500, "Command not found")
		return
	}
	if cmdObj.RequireParam() && param == "" {
		Conn.writeMessage(553, "action aborted, required param missing")
	} else if cmdObj.RequireAuth() && Conn.user == "" {
		Conn.writeMessage(530, "not logged in")
	} else {
		cmdObj.Execute(Conn, param)
	}
}

func (Conn *Conn) parseLine(line string) (string, string) {
	params := strings.SplitN(strings.Trim(line, "\r\n"), " ", 2)
	if len(params) == 1 {
		return params[0], ""
	}
	return params[0], strings.TrimSpace(params[1])
}

// writeMessage will send a standard FTP response back to the client.
func (Conn *Conn) writeMessage(code int, message string) (wrote int, err error) {
	Conn.logger.PrintResponse(code, message)
	line := fmt.Sprintf("%d %s\r\n", code, message)
	wrote, err = Conn.controlWriter.WriteString(line)
	Conn.controlWriter.Flush()
	return
}

// buildPath takes a client supplied path or filename and generates a safe
// absolute path within their account sandbox.
//
//    buildpath("/")
//    => "/"
//    buildpath("one.txt")
//    => "/one.txt"
//    buildpath("/files/two.txt")
//    => "/files/two.txt"
//    buildpath("files/two.txt")
//    => "files/two.txt"
//    buildpath("/../../../../etc/passwd")
//    => "/etc/passwd"
//
// The driver implementation is responsible for deciding how to treat this path.
// Obviously they MUST NOT just read the path off disk. The probably want to
// prefix the path with something to scope the users access to a sandbox.
func (Conn *Conn) buildPath(filename string) (fullPath string) {
	if len(filename) > 0 && filename[0:1] == "/" {
		fullPath = filepath.Clean(filename)
	} else if len(filename) > 0 && filename != "-a" {
		fullPath = filepath.Clean(Conn.namePrefix + "/" + filename)
	} else {
		fullPath = filepath.Clean(Conn.namePrefix)
	}
	fullPath = strings.Replace(fullPath, "//", "/", -1)
	return
}

// sendOutofbandData will send a string to the client via the currently open
// data socket. Assumes the socket is open and ready to be used.
func (Conn *Conn) sendOutofbandData(data []byte) {
	bytes := len(data)
	if Conn.dataConn != nil {
		Conn.dataConn.Write(data)
		Conn.dataConn.Close()
		Conn.dataConn = nil
	}
	message := "Closing data connection, sent " + strconv.Itoa(bytes) + " bytes"
	Conn.writeMessage(226, message)
}

func (Conn *Conn) sendOutofBandDataWriter(data io.ReadCloser) error {
	Conn.lastFilePos = 0
	bytes, err := io.Copy(Conn.dataConn, data)
	if err != nil {
		Conn.dataConn.Close()
		Conn.dataConn = nil
		return err
	}
	message := "Closing data connection, sent " + strconv.Itoa(int(bytes)) + " bytes"
	Conn.writeMessage(226, message)
	Conn.dataConn.Close()
	Conn.dataConn = nil

	return nil
}
<commit_msg>Close the connection on the server when the client closes it first<commit_after>package server

import (
	"bufio"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	defaultWelcomeMessage = "Welcome to the Go FTP Server"
)

type Conn struct {
	conn          net.Conn
	controlReader *bufio.Reader
	controlWriter *bufio.Writer
	dataConn      DataSocket
	driver        Driver
	auth          Auth
	logger        *Logger
	server        *Server
	sessionId     string
	namePrefix    string
	reqUser       string
	user          string
	renameFrom    string
	lastFilePos   int64
	appendData    bool
}

func (conn *Conn) LoginUser() string {
	return conn.user
}

func (conn *Conn) IsLogin() bool {
	return len(conn.user) > 0
}

// returns a random 20 char string that can be used as a unique session ID
func newSessionId() string {
	hash := sha256.New()
	_, err := io.CopyN(hash, rand.Reader, 50)
	if err != nil {
		return "????????????????????"
	}
	md := hash.Sum(nil)
	mdStr := hex.EncodeToString(md)
	return mdStr[0:20]
}

// Serve starts an endless loop that reads FTP commands from the client and
// responds appropriately. terminated is a channel that will receive a true
// message when the connection closes. This loop will be running inside a
// goroutine, so use this channel to be notified when the connection can be
// cleaned up.
func (conn *Conn) Serve() {
	conn.logger.Print("Connection Established")
	// send welcome
	conn.writeMessage(220, conn.server.WelcomeMessage)
	// read commands
	for {
		line, err := conn.controlReader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}

			conn.logger.Print(fmt.Sprintln("read error:", err))
			break
		}
		conn.receiveLine(line)
	}
	conn.Close()
	conn.logger.Print("Connection Terminated")
}

// Close will manually close this connection, even if the client isn't ready.
func (conn *Conn) Close() {
	conn.conn.Close()
	if conn.dataConn != nil {
		conn.dataConn.Close()
		conn.dataConn = nil
	}
}

// receiveLine accepts a single line FTP command and co-ordinates an
// appropriate response.
func (Conn *Conn) receiveLine(line string) {
	command, param := Conn.parseLine(line)
	Conn.logger.PrintCommand(command, param)
	cmdObj := commands[strings.ToUpper(command)]
	if cmdObj == nil {
		Conn.writeMessage(500, "Command not found")
		return
	}
	if cmdObj.RequireParam() && param == "" {
		Conn.writeMessage(553, "action aborted, required param missing")
	} else if cmdObj.RequireAuth() && Conn.user == "" {
		Conn.writeMessage(530, "not logged in")
	} else {
		cmdObj.Execute(Conn, param)
	}
}

func (Conn *Conn) parseLine(line string) (string, string) {
	params := strings.SplitN(strings.Trim(line, "\r\n"), " ", 2)
	if len(params) == 1 {
		return params[0], ""
	}
	return params[0], strings.TrimSpace(params[1])
}

// writeMessage will send a standard FTP response back to the client.
func (Conn *Conn) writeMessage(code int, message string) (wrote int, err error) {
	Conn.logger.PrintResponse(code, message)
	line := fmt.Sprintf("%d %s\r\n", code, message)
	wrote, err = Conn.controlWriter.WriteString(line)
	Conn.controlWriter.Flush()
	return
}

// buildPath takes a client supplied path or filename and generates a safe
// absolute path within their account sandbox.
//
//    buildpath("/")
//    => "/"
//    buildpath("one.txt")
//    => "/one.txt"
//    buildpath("/files/two.txt")
//    => "/files/two.txt"
//    buildpath("files/two.txt")
//    => "files/two.txt"
//    buildpath("/../../../../etc/passwd")
//    => "/etc/passwd"
//
// The driver implementation is responsible for deciding how to treat this path.
// Obviously they MUST NOT just read the path off disk. The probably want to
// prefix the path with something to scope the users access to a sandbox.
func (Conn *Conn) buildPath(filename string) (fullPath string) {
	if len(filename) > 0 && filename[0:1] == "/" {
		fullPath = filepath.Clean(filename)
	} else if len(filename) > 0 && filename != "-a" {
		fullPath = filepath.Clean(Conn.namePrefix + "/" + filename)
	} else {
		fullPath = filepath.Clean(Conn.namePrefix)
	}
	fullPath = strings.Replace(fullPath, "//", "/", -1)
	return
}

// sendOutofbandData will send a string to the client via the currently open
// data socket. Assumes the socket is open and ready to be used.
func (Conn *Conn) sendOutofbandData(data []byte) {
	bytes := len(data)
	if Conn.dataConn != nil {
		Conn.dataConn.Write(data)
		Conn.dataConn.Close()
		Conn.dataConn = nil
	}
	message := "Closing data connection, sent " + strconv.Itoa(bytes) + " bytes"
	Conn.writeMessage(226, message)
}

func (Conn *Conn) sendOutofBandDataWriter(data io.ReadCloser) error {
	Conn.lastFilePos = 0
	bytes, err := io.Copy(Conn.dataConn, data)
	if err != nil {
		Conn.dataConn.Close()
		Conn.dataConn = nil
		return err
	}
	message := "Closing data connection, sent " + strconv.Itoa(int(bytes)) + " bytes"
	Conn.writeMessage(226, message)
	Conn.dataConn.Close()
	Conn.dataConn = nil

	return nil
}
<|endoftext|>


<commit_before>package neptulon

import (
	"log"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"

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

// Conn is a client connection.
type Conn struct {
	ID         string
	Session    *cmap.CMap
	middleware []func(ctx *ReqCtx) error
	resRoutes  *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws         *websocket.Conn
	deadline   time.Duration
}

// NewConn creates a new Neptulon connection wrapping given websocket.Conn.
func NewConn(ws *websocket.Conn, middleware []func(ctx *ReqCtx) error) (*Conn, error) {
	id, err := shortid.UUID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:         id,
		Session:    cmap.New(),
		middleware: middleware,
		resRoutes:  cmap.New(),
		ws:         ws,
		deadline:   time.Second * time.Duration(300),
	}, nil
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
func (c *Conn) SetDeadline(seconds int) {
	c.deadline = time.Second * time.Duration(seconds)
}

// StartReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) StartReceive() {
	// append the last middleware to request stack, which will write the response to connection, if any
	c.middleware = append(c.middleware, func(ctx *ReqCtx) error {
		if ctx.Res != nil || ctx.Err != nil {
			return ctx.Conn.sendResponse(ctx.ID, ctx.Res, ctx.Err)
		}

		return nil
	})

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			log.Println("Error while receiving message:", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
				log.Println("Error while handling request:", err)
				break
			}

			continue
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
			c.resRoutes.Delete(m.ID)
			if err != nil {
				log.Println("Error while handling response:", err)
				break
			}
		} else {
			log.Println("Error while handling response: got response to a request with unknown ID:", m.ID)
			break
		}
	}
}

// SendRequest sends a JSON-RPC request through the connection with an auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequest(method string, params interface{}, resHandler func(res *Response) error) (reqID string, err error) {
	id, err := shortid.UUID()
	if err != nil {
		return "", err
	}

	req := Request{ID: id, Method: method, Params: params}
	if err = c.send(req); err != nil {
		return "", err
	}

	c.resRoutes.Set(req.ID, resHandler)
	return id, nil
}

// SendRequestArr sends a JSON-RPC request through the connection, with array params and auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequestArr(method string, resHandler func(res *Response) error, params ...interface{}) (reqID string, err error) {
	return c.SendRequest(method, params, resHandler)
}

// Close closes a connection.
func (c *Conn) Close() error {
	return c.ws.Close()
}

// SendResponse sends a JSON-RPC response message through the connection.
func (c *Conn) sendResponse(id string, result interface{}, err *ResError) error {
	return c.send(Response{ID: id, Result: result, Error: err})
}

// Send sends the given message through the connection.
func (c *Conn) send(msg interface{}) error {
	if err := c.ws.SetWriteDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Send(c.ws, msg)
}

// Receive receives message from the connection.
func (c *Conn) receive(msg *message) error {
	if err := c.ws.SetReadDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Receive(c.ws, &msg)
}
<commit_msg>add Conn.Middleware<commit_after>package neptulon

import (
	"log"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"

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

// Conn is a client connection.
type Conn struct {
	ID         string
	Session    *cmap.CMap
	middleware []func(ctx *ReqCtx) error
	resRoutes  *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws         *websocket.Conn
	deadline   time.Duration
}

// NewConn creates a new Neptulon connection wrapping given websocket.Conn.
func NewConn(ws *websocket.Conn, middleware []func(ctx *ReqCtx) error) (*Conn, error) {
	id, err := shortid.UUID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:         id,
		Session:    cmap.New(),
		middleware: middleware,
		resRoutes:  cmap.New(),
		ws:         ws,
		deadline:   time.Second * time.Duration(300),
	}, nil
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
func (c *Conn) SetDeadline(seconds int) {
	c.deadline = time.Second * time.Duration(seconds)
}

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...func(ctx *ReqCtx) error) {
	c.middleware = append(c.middleware, middleware...)
}

// StartReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) StartReceive() {
	// append the last middleware to request stack, which will write the response to connection, if any
	c.middleware = append(c.middleware, func(ctx *ReqCtx) error {
		if ctx.Res != nil || ctx.Err != nil {
			return ctx.Conn.sendResponse(ctx.ID, ctx.Res, ctx.Err)
		}

		return nil
	})

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			log.Println("Error while receiving message:", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
				log.Println("Error while handling request:", err)
				break
			}

			continue
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
			c.resRoutes.Delete(m.ID)
			if err != nil {
				log.Println("Error while handling response:", err)
				break
			}
		} else {
			log.Println("Error while handling response: got response to a request with unknown ID:", m.ID)
			break
		}
	}
}

// SendRequest sends a JSON-RPC request through the connection with an auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequest(method string, params interface{}, resHandler func(res *Response) error) (reqID string, err error) {
	id, err := shortid.UUID()
	if err != nil {
		return "", err
	}

	req := Request{ID: id, Method: method, Params: params}
	if err = c.send(req); err != nil {
		return "", err
	}

	c.resRoutes.Set(req.ID, resHandler)
	return id, nil
}

// SendRequestArr sends a JSON-RPC request through the connection, with array params and auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequestArr(method string, resHandler func(res *Response) error, params ...interface{}) (reqID string, err error) {
	return c.SendRequest(method, params, resHandler)
}

// Close closes a connection.
func (c *Conn) Close() error {
	return c.ws.Close()
}

// SendResponse sends a JSON-RPC response message through the connection.
func (c *Conn) sendResponse(id string, result interface{}, err *ResError) error {
	return c.send(Response{ID: id, Result: result, Error: err})
}

// Send sends the given message through the connection.
func (c *Conn) send(msg interface{}) error {
	if err := c.ws.SetWriteDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Send(c.ws, msg)
}

// Receive receives message from the connection.
func (c *Conn) receive(msg *message) error {
	if err := c.ws.SetReadDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Receive(c.ws, &msg)
}
<|endoftext|>


<commit_before>// Copyright © 2015 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package core

import (
	"fmt"
	. "github.com/thethingsnetwork/core/lorawan/semtech"
)

type BrokerAddress string
type GatewayAddress string

type Router interface {
	// HandlerError manages all kind of error that occur during the router lifecycle
	HandleError(err interface{})

	// HandleUplink manages uplink packets coming from a gateway
	HandleUplink(packet Packet, gateway GatewayAddress)

	// HandleDownlink manages downlink packets coming from a broker
	HandleDownlink(payload Payload, broker BrokerAddress)

	// RegisterDevice associates a device address to a set of brokers for a given period
	RegisterDevice(devAddr DeviceAddress, broAddrs ...BrokerAddress)

	// Connect the router to its adapters
	Connect(upAdapter GatewayRouterAdapter, downAdapter RouterBrokerAdapter)
}

// The error types belows are going to be more complex in order to handle custom behavior for
// each error type.
var ErrBadOptions error = fmt.Errorf("Unreckonized or invalid options")
var ErrBadGatewayAddress error = fmt.Errorf("Invalid gateway address")
var ErrMissingConnection error = fmt.Errorf("Can't proceed without establishing connection")
var ErrInvalidPacket error = fmt.Errorf("Invalid semtech packet")

type Adapter interface {
	// Establish the adapter connection, whatever protocol is being used.
	Listen(router Router, options interface{}) error
}

type GatewayRouterAdapter interface {
	Adapter
	// Ack allows the router to send back a response to a gateway. The name of the method is quite a
	// bad call and will probably change soon.
	Ack(router Router, packet Packet, gateway GatewayAddress)
}

type ErrDownlink error
type ErrForward error
type ErrBroadcast error
type RouterBrokerAdapter interface {
	Adapter

	// Broadcast makes the adapter discover all available brokers by sending them a the given packets.
	//
	// We assume that broadcast is also registering a device address towards the router depending
	// on the brokers responses.
	Broadcast(router Router, payload Payload, broAddrs ...BrokerAddress)

	// Forward is an explicit forwarding of a packet which is known being handled by a set of
	// brokers. None of the contacted broker is supposed to reject the incoming payload; They all
	// ave been queried before and are known as dedicated brokers for the related end-device.
	Forward(router Router, payload Payload, broAddrs ...BrokerAddress)
}
<commit_msg>[router] Change core signatures -> prefer early and direct error handling<commit_after>// Copyright © 2015 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package core

import (
	"fmt"
	. "github.com/thethingsnetwork/core/lorawan/semtech"
)

type BrokerAddress string
type GatewayAddress string

type Router interface {
	// HandlerError manages all kind of error that occur during the router lifecycle
	HandleError(err interface{})

	// HandleUplink manages uplink packets coming from a gateway
	HandleUplink(packet Packet, gateway GatewayAddress)

	// HandleDownlink manages downlink packets coming from a broker
	HandleDownlink(payload Payload, broker BrokerAddress)

	// RegisterDevice associates a device address to a set of brokers for a given period
	RegisterDevice(devAddr DeviceAddress, broAddrs ...BrokerAddress)

	// Connect the router to its adapters
	Connect(upAdapter GatewayRouterAdapter, downAdapter RouterBrokerAdapter)
}

// The error types belows are going to be more complex in order to handle custom behavior for
// each error type.
var ErrBadOptions error = fmt.Errorf("Unreckonized or invalid options")
var ErrBadGatewayAddress error = fmt.Errorf("Invalid gateway address")
var ErrMissingConnection error = fmt.Errorf("Can't proceed without establishing connection")
var ErrInvalidPacket error = fmt.Errorf("Invalid semtech packet")

type Adapter interface {
	// Establish the adapter connection, whatever protocol is being used.
	Listen(router Router, options interface{}) error
}

type GatewayRouterAdapter interface {
	Adapter
	// Ack allows the router to send back a response to a gateway. The name of the method is quite a
	// bad call and will probably change soon.
	Ack(router Router, packet Packet, gateway GatewayAddress) error
}

type ErrDownlink error
type ErrForward error
type ErrBroadcast error
type RouterBrokerAdapter interface {
	Adapter

	// Broadcast makes the adapter discover all available brokers by sending them a the given packets.
	//
	// We assume that broadcast is also registering a device address towards the router depending
	// on the brokers responses.
	Broadcast(router Router, payload Payload, broAddrs ...BrokerAddress) error

	// Forward is an explicit forwarding of a packet which is known being handled by a set of
	// brokers. None of the contacted broker is supposed to reject the incoming payload; They all
	// ave been queried before and are known as dedicated brokers for the related end-device.
	Forward(router Router, payload Payload, broAddrs ...BrokerAddress) error
}
<|endoftext|>


<commit_before>package db

import (
    "fmt"
    "regexp"
    "errors"
    "strings"
)

var (
	// unquoted array values must not contain: (" , \ { } whitespace NULL)
	// and must be at least one char
	unquotedChar  = `[^",\\{}\s(NULL)]`
	unquotedValue = fmt.Sprintf("(%s)+", unquotedChar)

	// quoted array values are surrounded by double quotes, can be any
	// character except " or \, which must be backslash escaped:
	quotedChar  = `[^"\\]|\\"|\\\\`
	quotedValue = fmt.Sprintf("\"(%s)*\"", quotedChar)

	// an array value may be either quoted or unquoted:
	arrayValue = fmt.Sprintf("(?P<value>(%s|%s))", unquotedValue, quotedValue)

	// Array values are separated with a comma IF there is more than one value:
	arrayExp = regexp.MustCompile(fmt.Sprintf("((%s)(,)?)", arrayValue))

	valueIndex int
)

type StringSlice []string

// Implements sql.Scanner for the String slice type
// Scanners take the database value (in this case as a byte slice)
// and sets the value of the type.  Here we cast to a string and
// do a regexp based parse
func (s *StringSlice) Scan(src interface{}) error {
	asBytes, ok := src.([]byte)
	if !ok {
		return error(errors.New("Scan source was not []bytes"))
	}
	asString := string(asBytes)
	parsed := parseArray(asString)
	(*s) = StringSlice(parsed)
	return nil
}

func parseArray(array string) []string {
	results := make([]string, 0)
	matches := arrayExp.FindAllStringSubmatch(array, -1)
	for _, match := range matches {
		s := match[valueIndex]
		// the string _might_ be wrapped in quotes, so trim them:
		s = strings.Trim(s, "\"")
		results = append(results, s)
	}
	return results
}
<commit_msg>go fmt<commit_after>package db

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
)

var (
	// unquoted array values must not contain: (" , \ { } whitespace NULL)
	// and must be at least one char
	unquotedChar  = `[^",\\{}\s(NULL)]`
	unquotedValue = fmt.Sprintf("(%s)+", unquotedChar)

	// quoted array values are surrounded by double quotes, can be any
	// character except " or \, which must be backslash escaped:
	quotedChar  = `[^"\\]|\\"|\\\\`
	quotedValue = fmt.Sprintf("\"(%s)*\"", quotedChar)

	// an array value may be either quoted or unquoted:
	arrayValue = fmt.Sprintf("(?P<value>(%s|%s))", unquotedValue, quotedValue)

	// Array values are separated with a comma IF there is more than one value:
	arrayExp = regexp.MustCompile(fmt.Sprintf("((%s)(,)?)", arrayValue))

	valueIndex int
)

type StringSlice []string

// Implements sql.Scanner for the String slice type
// Scanners take the database value (in this case as a byte slice)
// and sets the value of the type.  Here we cast to a string and
// do a regexp based parse
func (s *StringSlice) Scan(src interface{}) error {
	asBytes, ok := src.([]byte)
	if !ok {
		return error(errors.New("Scan source was not []bytes"))
	}
	asString := string(asBytes)
	parsed := parseArray(asString)
	(*s) = StringSlice(parsed)
	return nil
}

func parseArray(array string) []string {
	results := make([]string, 0)
	matches := arrayExp.FindAllStringSubmatch(array, -1)
	for _, match := range matches {
		s := match[valueIndex]
		// the string _might_ be wrapped in quotes, so trim them:
		s = strings.Trim(s, "\"")
		results = append(results, s)
	}
	return results
}
<|endoftext|>


<commit_before>// Package bundle manages translations for multiple languages.
package bundle

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"reflect"

	//	"launchpad.net/goyaml"

	"path/filepath"

	"github.com/nicksnyder/go-i18n/i18n/language"
	"github.com/nicksnyder/go-i18n/i18n/translation"
)

// TranslateFunc is a copy of i18n.TranslateFunc to avoid a circular dependency.
type TranslateFunc func(translationID string, args ...interface{}) string

// Bundle stores the translations for multiple languages.
type Bundle struct {
	// The primary translations for a language tag and translation id.
	translations map[string]map[string]translation.Translation

	// Translations that can be used when an exact language match is not possible.
	fallbackTranslations map[string]map[string]translation.Translation
}

// New returns an empty bundle.
func New() *Bundle {
	return &Bundle{
		translations:         make(map[string]map[string]translation.Translation),
		fallbackTranslations: make(map[string]map[string]translation.Translation),
	}
}

// MustLoadTranslationFile is similar to LoadTranslationFile
// except it panics if an error happens.
func (b *Bundle) MustLoadTranslationFile(filename string) {
	if err := b.LoadTranslationFile(filename); err != nil {
		panic(err)
	}
}

// LoadTranslationFile loads the translations from filename into memory.
//
// The language that the translations are associated with is parsed from the filename (e.g. en-US.json).
//
// Generally you should load translation files once during your program's initialization.
func (b *Bundle) LoadTranslationFile(filename string) error {
	buf, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}
	return b.ParseTranslationFileBytes(filename, buf)
}

// ParseTranslationFileBytes is similar to LoadTranslationFile except it parses the bytes in buf.
//
// It is useful for parsing translation files embedded with go-bindata.
func (b *Bundle) ParseTranslationFileBytes(filename string, buf []byte) error {
	basename := filepath.Base(filename)
	langs := language.Parse(basename)
	switch l := len(langs); {
	case l == 0:
		return fmt.Errorf("no language found in %q", basename)
	case l > 1:
		return fmt.Errorf("multiple languages found in filename %q: %v; expected one", basename, langs)
	}
	translations, err := parseTranslations(filename, buf)
	if err != nil {
		return err
	}
	b.AddTranslation(langs[0], translations...)
	return nil
}

func parseTranslations(filename string, buf []byte) ([]translation.Translation, error) {
	var unmarshalFunc func([]byte, interface{}) error
	switch format := filepath.Ext(filename); format {
	case ".json":
		unmarshalFunc = json.Unmarshal
	/*
		case ".yaml":
			unmarshalFunc = goyaml.Unmarshal
	*/
	default:
		return nil, fmt.Errorf("unsupported file extension %s", format)
	}

	var translationsData []map[string]interface{}
	if len(buf) > 0 {
		if err := unmarshalFunc(buf, &translationsData); err != nil {
			return nil, err
		}
	}

	translations := make([]translation.Translation, 0, len(translationsData))
	for i, translationData := range translationsData {
		t, err := translation.NewTranslation(translationData)
		if err != nil {
			return nil, fmt.Errorf("unable to parse translation #%d in %s because %s\n%v", i, filename, err, translationData)
		}
		translations = append(translations, t)
	}
	return translations, nil
}

// AddTranslation adds translations for a language.
//
// It is useful if your translations are in a format not supported by LoadTranslationFile.
func (b *Bundle) AddTranslation(lang *language.Language, translations ...translation.Translation) {
	if b.translations[lang.Tag] == nil {
		b.translations[lang.Tag] = make(map[string]translation.Translation, len(translations))
	}
	currentTranslations := b.translations[lang.Tag]
	for _, newTranslation := range translations {
		if currentTranslation := currentTranslations[newTranslation.ID()]; currentTranslation != nil {
			currentTranslations[newTranslation.ID()] = currentTranslation.Merge(newTranslation)
		} else {
			currentTranslations[newTranslation.ID()] = newTranslation
		}
	}

	// lang can provide translations for less specific language tags.
	for _, tag := range lang.MatchingTags() {
		b.fallbackTranslations[tag] = currentTranslations
	}
}

// Translations returns all translations in the bundle.
func (b *Bundle) Translations() map[string]map[string]translation.Translation {
	return b.translations
}

// MustTfunc is similar to Tfunc except it panics if an error happens.
func (b *Bundle) MustTfunc(pref string, prefs ...string) TranslateFunc {
	tfunc, err := b.Tfunc(pref, prefs...)
	if err != nil {
		panic(err)
	}
	return tfunc
}

// MustTfuncAndLanguage is similar to TfuncAndLanguage except it panics if an error happens.
func (b *Bundle) MustTfuncAndLanguage(pref string, prefs ...string) (TranslateFunc, *language.Language) {
	tfunc, language, err := b.TfuncAndLanguage(pref, prefs...)
	if err != nil {
		panic(err)
	}
	return tfunc, language
}

// Tfunc is similar to TfuncAndLanguage except is doesn't return the Language.
func (b *Bundle) Tfunc(pref string, prefs ...string) (TranslateFunc, error) {
	tfunc, _, err := b.TfuncAndLanguage(pref, prefs...)
	return tfunc, err
}

// TfuncAndLanguage returns a TranslateFunc for the first Language that
// has a non-zero number of translations in the bundle.
//
// The returned Language matches the the first language preference that could be satisfied,
// but this may not strictly match the language of the translations used to satisfy that preference.
//
// For example, the user may request "zh". If there are no translations for "zh" but there are translations
// for "zh-cn", then the translations for "zh-cn" will be used but the returned Language will be "zh".
//
// It can parse languages from Accept-Language headers (RFC 2616),
// but it assumes weights are monotonically decreasing.
func (b *Bundle) TfuncAndLanguage(pref string, prefs ...string) (TranslateFunc, *language.Language, error) {
	lang := b.supportedLanguage(pref, prefs...)
	var err error
	if lang == nil {
		err = fmt.Errorf("no supported languages found %#v", append(prefs, pref))
	}
	return func(translationID string, args ...interface{}) string {
		return b.translate(lang, translationID, args...)
	}, lang, err
}

// supportedLanguage returns the first language which
// has a non-zero number of translations in the bundle.
func (b *Bundle) supportedLanguage(pref string, prefs ...string) *language.Language {
	lang := b.translatedLanguage(pref)
	if lang == nil {
		for _, pref := range prefs {
			lang = b.translatedLanguage(pref)
			if lang != nil {
				break
			}
		}
	}
	return lang
}

func (b *Bundle) translatedLanguage(src string) *language.Language {
	langs := language.Parse(src)
	for _, lang := range langs {
		if len(b.translations[lang.Tag]) > 0 ||
			len(b.fallbackTranslations[lang.Tag]) > 0 {
			return lang
		}
	}
	return nil
}

func (b *Bundle) translate(lang *language.Language, translationID string, args ...interface{}) string {
	if lang == nil {
		return translationID
	}

	translations := b.translations[lang.Tag]
	if translations == nil {
		translations = b.fallbackTranslations[lang.Tag]
		if translations == nil {
			return translationID
		}
	}

	translation := translations[translationID]
	if translation == nil {
		return translationID
	}

	var count interface{}
	if len(args) > 0 && isNumber(args[0]) {
		count = args[0]
		args = args[1:]
	}

	plural, _ := lang.Plural(count)
	template := translation.Template(plural)
	if template == nil {
		return translationID
	}

	var data map[string]interface{}
	if len(args) > 0 {
		data = toMap(args[0])
	}

	if isNumber(count) {
		if data == nil {
			data = map[string]interface{}{"Count": count}
		} else {
			data["Count"] = count
		}
	}

	s := template.Execute(data)
	if s == "" {
		return translationID
	}
	return s
}

func isNumber(n interface{}) bool {
	switch n.(type) {
	case int, int8, int16, int32, int64, string:
		return true
	}
	return false
}

func toMap(input interface{}) map[string]interface{} {
	v := reflect.ValueOf(input)
	switch v.Kind() {
	case reflect.Map:
		data, _ := input.(map[string]interface{})
		return data
	case reflect.Ptr:
		return structToMap(v.Elem())
	case reflect.Struct:
		return structToMap(v)
	default:
		panic(fmt.Errorf("i18n: cannot handle type %s", v.Kind()))
	}
}
<commit_msg>bundle: toMap() should recurse if it encounters a pointer<commit_after>// Package bundle manages translations for multiple languages.
package bundle

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"reflect"

	//	"launchpad.net/goyaml"

	"path/filepath"

	"github.com/nicksnyder/go-i18n/i18n/language"
	"github.com/nicksnyder/go-i18n/i18n/translation"
)

// TranslateFunc is a copy of i18n.TranslateFunc to avoid a circular dependency.
type TranslateFunc func(translationID string, args ...interface{}) string

// Bundle stores the translations for multiple languages.
type Bundle struct {
	// The primary translations for a language tag and translation id.
	translations map[string]map[string]translation.Translation

	// Translations that can be used when an exact language match is not possible.
	fallbackTranslations map[string]map[string]translation.Translation
}

// New returns an empty bundle.
func New() *Bundle {
	return &Bundle{
		translations:         make(map[string]map[string]translation.Translation),
		fallbackTranslations: make(map[string]map[string]translation.Translation),
	}
}

// MustLoadTranslationFile is similar to LoadTranslationFile
// except it panics if an error happens.
func (b *Bundle) MustLoadTranslationFile(filename string) {
	if err := b.LoadTranslationFile(filename); err != nil {
		panic(err)
	}
}

// LoadTranslationFile loads the translations from filename into memory.
//
// The language that the translations are associated with is parsed from the filename (e.g. en-US.json).
//
// Generally you should load translation files once during your program's initialization.
func (b *Bundle) LoadTranslationFile(filename string) error {
	buf, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}
	return b.ParseTranslationFileBytes(filename, buf)
}

// ParseTranslationFileBytes is similar to LoadTranslationFile except it parses the bytes in buf.
//
// It is useful for parsing translation files embedded with go-bindata.
func (b *Bundle) ParseTranslationFileBytes(filename string, buf []byte) error {
	basename := filepath.Base(filename)
	langs := language.Parse(basename)
	switch l := len(langs); {
	case l == 0:
		return fmt.Errorf("no language found in %q", basename)
	case l > 1:
		return fmt.Errorf("multiple languages found in filename %q: %v; expected one", basename, langs)
	}
	translations, err := parseTranslations(filename, buf)
	if err != nil {
		return err
	}
	b.AddTranslation(langs[0], translations...)
	return nil
}

func parseTranslations(filename string, buf []byte) ([]translation.Translation, error) {
	var unmarshalFunc func([]byte, interface{}) error
	switch format := filepath.Ext(filename); format {
	case ".json":
		unmarshalFunc = json.Unmarshal
	/*
		case ".yaml":
			unmarshalFunc = goyaml.Unmarshal
	*/
	default:
		return nil, fmt.Errorf("unsupported file extension %s", format)
	}

	var translationsData []map[string]interface{}
	if len(buf) > 0 {
		if err := unmarshalFunc(buf, &translationsData); err != nil {
			return nil, err
		}
	}

	translations := make([]translation.Translation, 0, len(translationsData))
	for i, translationData := range translationsData {
		t, err := translation.NewTranslation(translationData)
		if err != nil {
			return nil, fmt.Errorf("unable to parse translation #%d in %s because %s\n%v", i, filename, err, translationData)
		}
		translations = append(translations, t)
	}
	return translations, nil
}

// AddTranslation adds translations for a language.
//
// It is useful if your translations are in a format not supported by LoadTranslationFile.
func (b *Bundle) AddTranslation(lang *language.Language, translations ...translation.Translation) {
	if b.translations[lang.Tag] == nil {
		b.translations[lang.Tag] = make(map[string]translation.Translation, len(translations))
	}
	currentTranslations := b.translations[lang.Tag]
	for _, newTranslation := range translations {
		if currentTranslation := currentTranslations[newTranslation.ID()]; currentTranslation != nil {
			currentTranslations[newTranslation.ID()] = currentTranslation.Merge(newTranslation)
		} else {
			currentTranslations[newTranslation.ID()] = newTranslation
		}
	}

	// lang can provide translations for less specific language tags.
	for _, tag := range lang.MatchingTags() {
		b.fallbackTranslations[tag] = currentTranslations
	}
}

// Translations returns all translations in the bundle.
func (b *Bundle) Translations() map[string]map[string]translation.Translation {
	return b.translations
}

// MustTfunc is similar to Tfunc except it panics if an error happens.
func (b *Bundle) MustTfunc(pref string, prefs ...string) TranslateFunc {
	tfunc, err := b.Tfunc(pref, prefs...)
	if err != nil {
		panic(err)
	}
	return tfunc
}

// MustTfuncAndLanguage is similar to TfuncAndLanguage except it panics if an error happens.
func (b *Bundle) MustTfuncAndLanguage(pref string, prefs ...string) (TranslateFunc, *language.Language) {
	tfunc, language, err := b.TfuncAndLanguage(pref, prefs...)
	if err != nil {
		panic(err)
	}
	return tfunc, language
}

// Tfunc is similar to TfuncAndLanguage except is doesn't return the Language.
func (b *Bundle) Tfunc(pref string, prefs ...string) (TranslateFunc, error) {
	tfunc, _, err := b.TfuncAndLanguage(pref, prefs...)
	return tfunc, err
}

// TfuncAndLanguage returns a TranslateFunc for the first Language that
// has a non-zero number of translations in the bundle.
//
// The returned Language matches the the first language preference that could be satisfied,
// but this may not strictly match the language of the translations used to satisfy that preference.
//
// For example, the user may request "zh". If there are no translations for "zh" but there are translations
// for "zh-cn", then the translations for "zh-cn" will be used but the returned Language will be "zh".
//
// It can parse languages from Accept-Language headers (RFC 2616),
// but it assumes weights are monotonically decreasing.
func (b *Bundle) TfuncAndLanguage(pref string, prefs ...string) (TranslateFunc, *language.Language, error) {
	lang := b.supportedLanguage(pref, prefs...)
	var err error
	if lang == nil {
		err = fmt.Errorf("no supported languages found %#v", append(prefs, pref))
	}
	return func(translationID string, args ...interface{}) string {
		return b.translate(lang, translationID, args...)
	}, lang, err
}

// supportedLanguage returns the first language which
// has a non-zero number of translations in the bundle.
func (b *Bundle) supportedLanguage(pref string, prefs ...string) *language.Language {
	lang := b.translatedLanguage(pref)
	if lang == nil {
		for _, pref := range prefs {
			lang = b.translatedLanguage(pref)
			if lang != nil {
				break
			}
		}
	}
	return lang
}

func (b *Bundle) translatedLanguage(src string) *language.Language {
	langs := language.Parse(src)
	for _, lang := range langs {
		if len(b.translations[lang.Tag]) > 0 ||
			len(b.fallbackTranslations[lang.Tag]) > 0 {
			return lang
		}
	}
	return nil
}

func (b *Bundle) translate(lang *language.Language, translationID string, args ...interface{}) string {
	if lang == nil {
		return translationID
	}

	translations := b.translations[lang.Tag]
	if translations == nil {
		translations = b.fallbackTranslations[lang.Tag]
		if translations == nil {
			return translationID
		}
	}

	translation := translations[translationID]
	if translation == nil {
		return translationID
	}

	var count interface{}
	if len(args) > 0 && isNumber(args[0]) {
		count = args[0]
		args = args[1:]
	}

	plural, _ := lang.Plural(count)
	template := translation.Template(plural)
	if template == nil {
		return translationID
	}

	var data map[string]interface{}
	if len(args) > 0 {
		data = toMap(args[0])
	}

	if isNumber(count) {
		if data == nil {
			data = map[string]interface{}{"Count": count}
		} else {
			data["Count"] = count
		}
	}

	s := template.Execute(data)
	if s == "" {
		return translationID
	}
	return s
}

func isNumber(n interface{}) bool {
	switch n.(type) {
	case int, int8, int16, int32, int64, string:
		return true
	}
	return false
}

func toMap(input interface{}) map[string]interface{} {
	v := reflect.ValueOf(input)
	switch v.Kind() {
	case reflect.Map:
		data, _ := input.(map[string]interface{})
		return data
	case reflect.Ptr:
		return toMap(v.Elem().Interface())
	case reflect.Struct:
		return structToMap(v)
	default:
		panic(fmt.Errorf("i18n: cannot handle type %s", v.Kind()))
	}
}
<|endoftext|>