<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kcache "k8s.io/kubernetes/pkg/client/unversioned/cache"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcontrollerFramework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"

  "k8s.io/kubernetes/pkg/api"
  "k8s.io/kubernetes/pkg/fields"
  "k8s.io/kubernetes/pkg/labels"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client

	// DNS domain name.
	domain string

	//Nodes Name / valid
	nodes map[string]bool

	//DNS IDS
	ids map[string][]string
}

func Newkube2consul() *kube2consul {
	var k kube2consul
	k.nodes = make(map[string]bool)
	k.ids = make(map[string][]string)

	return &k
}

func (ks *kube2consul) removeDNS(record string) error {
	glog.V(2).Infof("Removing %s from DNS", record)
	for _,id := range ks.ids[record] {
		ks.consulClient.Agent().ServiceDeregister(id)
	}
	ks.ids[record] = []string{}
	return nil
}

func (ks *kube2consul) addDNS(record string, service *kapi.Service) error {
	if strings.Contains(record, ".") {
		glog.V(1).Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if ClusterIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.V(1).Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
		for n,s := range ks.nodes {
			if s {
				newId := n+record + service.Spec.Ports[i].Name

				asr := &consulapi.AgentServiceRegistration{
					ID:			 newId,
					Name: 	 record + "-" + service.Spec.Ports[i].Name,
					Address: n,
					Port:    service.Spec.Ports[i].NodePort,
					Tags: []string{"Kube"},
				}

				glog.V(2).Infof("Setting DNS record: %v -> %v:%d\n", record, service.Spec.ClusterIP, service.Spec.Ports[i].Port)
				if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
					return err
				}
				ks.ids[record] = append(ks.ids[record], newId)
			}
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		overrides := &kclientcmd.ConfigOverrides{}
		overrides.ClusterInfo.Server = masterUrl                                     // might be "", but that is OK
		rules := &kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile} // might be "", but that is OK
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	//glog.Infof("Name String: %s  %s", service, namespace)
	return fmt.Sprintf("%s.%s", service, namespace)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to services.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.addDNS(s.Name, s); err != nil {
			glog.V(1).Infof("Failed to add service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	glog.Info("Service remove")
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.removeDNS(s.Name); err != nil {
			glog.V(1).Infof("Failed to remove service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) updateNode(oldObj, newObj interface{}) {
	if n, ok := oldObj.(*kapi.Node); ok {
		name := n.Name
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue

		ks.nodes[name] = ready
	}

	if n, ok := newObj.(*kapi.Node); ok {
		name := n.Name
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue

		ks.nodes[name] = ready
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: func(oldObj, newObj interface{}) {
				ks.newService(newObj)
			},
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
}

func watchForNodes(kubeClient *kclient.Client, ks *kube2consul) kcache.Store {
	store, serviceController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    func(newObj interface{}) {
				glog.Info("Adding node!")
			},
			DeleteFunc: func(oldObj interface{}) {
				glog.Info("Node Removed!!")
			},
			UpdateFunc: ks.updateNode,
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
	return store
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := Newkube2consul()

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	glog.Info(kubeClient.ServerVersion())
	glog.Info(kubeClient.Services(kapi.NamespaceAll).Get("sensu-core"))

	pods, err := kubeClient.Pods(api.NamespaceDefault).List(labels.Everything(), fields.Everything())
	if err != nil {
	  for pod := range pods.Items {
			glog.Info(pod)
		}
	}

	watchForServices(kubeClient, ks)
	watchForNodes(kubeClient, ks)
	glog.Info("Watchers running")
	select {}
}
<commit_msg>Supports removal of services and nodes on removal of either a node or service. Excess consul entries no longer exist<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main // import "github.com/jmccarty3/kube2consul"

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	kcache "k8s.io/kubernetes/pkg/client/unversioned/cache"
	kclientcmd "k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	kframework "k8s.io/kubernetes/pkg/controller/framework"
	kcontrollerFramework "k8s.io/kubernetes/pkg/controller/framework"
	kSelector "k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"

  "k8s.io/kubernetes/pkg/api"
  "k8s.io/kubernetes/pkg/fields"
  "k8s.io/kubernetes/pkg/labels"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type nodeInformation struct {
	name string

	address string

	ready bool
	// map[service] DNS IDs
	ids map[string][]string
}

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client

	// DNS domain name.
	domain string

	//Nodes Name / valid
	nodes map[string]nodeInformation

	//All Services.
	services map[string]*kapi.Service
}

func Newkube2consul() *kube2consul {
	var k kube2consul
	k.nodes = make(map[string]nodeInformation)
	k.services = make(map[string]*kapi.Service)

	return &k
}

func NewnodeInformation() *nodeInformation {
	var n nodeInformation
	n.ready = false
	n.ids = make(map[string][]string)

	return &n
}

func Contains(s []string, e string) bool {
	for _,i := range s {
		if i == e {
			return true
		}
	}
	return false
}

func (ks *kube2consul) removeDNS(recordID string) error {
	glog.Infof("Removing %s from DNS", recordID)
	ks.consulClient.Agent().ServiceDeregister(recordID)
	return nil
}

func (ks *kube2consul) createDNS(record string, service *kapi.Service, node *nodeInformation) error {
	if strings.Contains(record, ".") {
		glog.Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if ClusterIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
			newId := node.name+record + service.Spec.Ports[i].Name

			asr := &consulapi.AgentServiceRegistration{
				ID:			 newId,
				Name: 	 record + "-" + service.Spec.Ports[i].Name,
				Address: node.address,
				Port:    service.Spec.Ports[i].NodePort,
				Tags: []string{"Kube"},
			}

			if Contains(node.ids[record], newId) == false {
				glog.Infof("Setting DNS record: %v -> %v:%d\n", asr.Name, asr.Address, asr.Port)
				if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
					return err
				}

				node.ids[record] = append(node.ids[record], newId)
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		overrides := &kclientcmd.ConfigOverrides{}
		overrides.ClusterInfo.Server = masterUrl                                     // might be "", but that is OK
		rules := &kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile} // might be "", but that is OK
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	//glog.Infof("Name String: %s  %s", service, namespace)
	//return fmt.Sprintf("%s.%s", service, namespace)
	return fmt.Sprintf("%s", service)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

// Returns a cache.ListWatch that gets all changes to services.
func createNodeLW(kubeClient *kclient.Client) *kcache.ListWatch {
	return kcache.NewListWatchFromClient(kubeClient, "nodes", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		ks.services[name] = s

		glog.V(2).Info("Creating Service: ", name)
		//Add to all existing nodes
		for _,node := range ks.nodes {
			if node.ready {
				ks.createDNS(name, s, &node)
			}
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		//Remove service from node
		for _,node := range ks.nodes {
			for _,id := range node.ids[name] {
				ks.removeDNS(id)
			}
			delete(node.ids, name)
		}
		delete(ks.services,name)
	}
}

func (ks *kube2consul) updateNode(oldObj, newObj interface{}) {
	if n, ok := newObj.(*kapi.Node); ok {
		ready := n.Status.Conditions[0].Status == kapi.ConditionTrue
		nodeInfo := ks.nodes[n.Name]

		if nodeInfo.ready != ready {
			nodeInfo.ready = ready
			if ready {
				for serviceName, service := range ks.services {
					ks.createDNS(serviceName, service, &nodeInfo)
				}
			}	else {
				for _,serviceIDs := range nodeInfo.ids {
					for _,serviceID := range serviceIDs {
						ks.removeDNS(serviceID)
					}
				}
				nodeInfo.ids = make(map[string][]string) //Clear tha map
			}
		}
	}
}

func (ks *kube2consul) newNode(newObj interface{}) {
	if node, ok := newObj.(*kapi.Node); ok {
		if _, ok := ks.nodes[node.Name]; !ok {
			glog.Info("Adding Node: ", node.Name)
			var newNode = *NewnodeInformation()
			newNode.name = node.Name
			newNode.address = node.Status.Addresses[0].Address

			ks.nodes[node.Name] = newNode
		}
	}
}

func (ks *kube2consul) removeNode(oldObj interface{}) {
	if node, ok := oldObj.(*kapi.Node); ok {
		if info, ok := ks.nodes[node.Name]; ok {
			glog.Info("Removing Node: ", node.Name)
			for _,idSet := range info.ids {
				for _,id := range idSet {
					ks.removeDNS(id)
				}
			}

			delete(ks.nodes, node.Name)
		}
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = kframework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: func(oldObj, newObj interface{}) {
				ks.newService(newObj)
			},
		},
	)
	go serviceController.Run(util.NeverStop)
}

func watchForNodes(kubeClient *kclient.Client, ks *kube2consul) kcache.Store {
	store, serviceController := kframework.NewInformer(
		createNodeLW(kubeClient),
		&kapi.Node{},
		resyncPeriod,
		kframework.ResourceEventHandlerFuncs{
			AddFunc: ks.newNode,
			DeleteFunc: ks.removeNode,
			UpdateFunc: ks.updateNode,
		},
	)
	glog.Info("About to call run!")
	go serviceController.Run(util.NeverStop)
	return store
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := Newkube2consul()

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	glog.Info(kubeClient.ServerVersion())
	glog.Info(kubeClient.Services(kapi.NamespaceAll).Get("sensu-core"))

	pods, err := kubeClient.Pods(api.NamespaceDefault).List(labels.Everything(), fields.Everything())
	if err != nil {
	  for pod := range pods.Items {
			glog.Info(pod)
		}
	}

	watchForServices(kubeClient, ks)
	watchForNodes(kubeClient, ks)
	glog.Info("Watchers running")
	select {}
}
<|endoftext|>


<commit_before>// Package torrent provides a torrent client implementation.
package torrent

import (
	"encoding/hex"
	"errors"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"

	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/resume"
	"github.com/cenkalti/rain/storage"
	"github.com/cenkalti/rain/storage/filestorage"
	"github.com/cenkalti/rain/torrent/internal/acceptor"
	"github.com/cenkalti/rain/torrent/internal/addrlist"
	"github.com/cenkalti/rain/torrent/internal/allocator"
	"github.com/cenkalti/rain/torrent/internal/announcer"
	"github.com/cenkalti/rain/torrent/internal/bitfield"
	"github.com/cenkalti/rain/torrent/internal/handshaker/incominghandshaker"
	"github.com/cenkalti/rain/torrent/internal/handshaker/outgoinghandshaker"
	"github.com/cenkalti/rain/torrent/internal/infodownloader"
	"github.com/cenkalti/rain/torrent/internal/magnet"
	"github.com/cenkalti/rain/torrent/internal/metainfo"
	"github.com/cenkalti/rain/torrent/internal/mse"
	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/piece"
	"github.com/cenkalti/rain/torrent/internal/piecedownloader"
	"github.com/cenkalti/rain/torrent/internal/torrentdata"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/httptracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/udptracker"
	"github.com/cenkalti/rain/torrent/internal/verifier"
)

const (
	maxPeerDial            = 40
	maxPeerAccept          = 40
	parallelPieceDownloads = 25 // running downloads, snubbed and choked peers don't count
	parallelInfoDownloads  = 1  // running downloads, snubbed peers don't count
)

var (
	// http://www.bittorrent.org/beps/bep_0020.html
	peerIDPrefix = []byte("-RN" + clientversion.Version + "-")
)

// Torrent connects to peers and downloads files from swarm.
type Torrent struct {
	*downloadSpec

	// Unique peer ID is generated per downloader.
	peerID [20]byte

	// Data provides IO access to pieces in torrent.
	data *torrentdata.Data

	// Contains state about the pieces in torrent.
	pieces []*piece.Piece

	// Contains pieces in sorted order for piece selection function.
	sortedPieces []*piece.Piece

	// Peers are sent to this channel when they are disconnected.
	peerDisconnectedC chan *peer.Peer

	// All messages coming from peers are sent to this channel.
	messages chan peer.Message

	// We keep connected peers in this map after they complete handshake phase.
	peers         map[*peer.Peer]struct{}
	incomingPeers map[*peer.Peer]struct{}
	outgoingPeers map[*peer.Peer]struct{}
	peersSnubbed  map[*peer.Peer]struct{}

	// Active piece downloads are kept in this map.
	pieceDownloaders        map[*peer.Peer]*piecedownloader.PieceDownloader
	pieceDownloadersSnubbed map[*peer.Peer]*piecedownloader.PieceDownloader
	pieceDownloadersChoked  map[*peer.Peer]*piecedownloader.PieceDownloader

	// When piece downloaders detects that a peer has snubbed us, it will send a signal to this channel.
	snubbedPieceDownloaderC chan *piecedownloader.PieceDownloader
	snubbedInfoDownloaderC  chan *infodownloader.InfoDownloader

	// Active metadata downloads are kept in this map.
	infoDownloaders        map[*peer.Peer]*infodownloader.InfoDownloader
	infoDownloadersSnubbed map[*peer.Peer]*infodownloader.InfoDownloader

	// A peer is optimistically unchoked regardless of their download rate.
	optimisticUnchokedPeer *peer.Peer

	// This channel is closed once all pieces are downloaded and verified.
	completeC chan struct{}

	// If any unrecoverable error occurs, it will be sent to this channel and download will be stopped.
	errC chan error

	// When Stop() is called, it will close this channel to signal run() function to stop.
	closeC chan struct{}

	// This channel will be closed after run loop exists.
	doneC chan struct{}

	// These are the channels for sending a message to run() loop.
	statsCommandC       chan statsRequest       // Stats()
	startCommandC       chan struct{}           // Start()
	stopCommandC        chan struct{}           // Stop()
	notifyErrorCommandC chan notifyErrorCommand // NotifyError()

	// Trackers send announce responses to this channel.
	addrsFromTrackers chan []*net.TCPAddr

	// Keeps a list of peer addresses to connect.
	addrList *addrlist.AddrList

	// New raw connections created by OutgoingHandshaker are sent to here.
	incomingConnC chan net.Conn

	// Keep a set of peer IDs to block duplicate connections.
	peerIDs map[[20]byte]struct{}

	// Listens for incoming peer connections.
	acceptor *acceptor.Acceptor

	// Special hash of info hash for encypted connection handshake.
	sKeyHash [20]byte

	// Tracker implementations for giving to announcers.
	trackersInstances []tracker.Tracker

	// Announces the status of torrent to trackers to get peer addresses.
	announcers []*announcer.Announcer

	// List of peers in handshake state.
	incomingHandshakers map[string]*incominghandshaker.IncomingHandshaker
	outgoingHandshakers map[string]*outgoinghandshaker.OutgoingHandshaker

	// Handshake results are sent to these channels by handshakers.
	incomingHandshakerResultC chan incominghandshaker.Result
	outgoingHandshakerResultC chan outgoinghandshaker.Result

	// When metadata of the torrent downloaded completely, a message is sent to this channel.
	infoDownloaderDoneC chan *infodownloader.InfoDownloader

	// When a piece is downloaded completely a message is sent to this channel.
	pieceDownloaderDoneC chan *piecedownloader.PieceDownloader

	// Announcers send a request to this channel to get information about the torrent.
	announcerRequestC chan *announcer.Request

	// A timer that ticks periodically to keep a certain number of peers unchoked.
	unchokeTimer  *time.Ticker
	unchokeTimerC <-chan time.Time

	// A timer that ticks periodically to keep a random peer unchoked regardless of its upload rate.
	optimisticUnchokeTimer  *time.Ticker
	optimisticUnchokeTimerC <-chan time.Time

	allocator          *allocator.Allocator
	allocatorProgressC chan allocator.Progress
	allocatorResultC   chan allocator.Result

	verifier          *verifier.Verifier
	verifierProgressC chan verifier.Progress
	verifierResultC   chan verifier.Result

	bytesDownloaded int64
	bytesWasted     int64

	log logger.Logger
}

// New returns a new torrent by reading a torrent metainfo file.
func New(r io.Reader, port int, sto storage.Storage) (*Torrent, error) {
	m, err := metainfo.New(r)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.Info.Hash,
		trackers: m.GetTrackers(),
		name:     m.Info.Name,
		port:     port,
		storage:  sto,
		info:     m.Info,
	}
	return newTorrent(spec)
}

// NewMagnet returns a new torrent by parsing a magnet link.
func NewMagnet(link string, port int, sto storage.Storage) (*Torrent, error) {
	m, err := magnet.New(link)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.InfoHash,
		trackers: m.Trackers,
		name:     m.Name,
		port:     port,
		storage:  sto,
	}
	return newTorrent(spec)
}

// NewResume returns a new torrent by loading all info from a resume.DB.
func NewResume(res resume.DB) (*Torrent, error) {
	spec, err := res.Read()
	if err != nil {
		return nil, err
	}
	if spec == nil {
		return nil, errors.New("no resume info")
	}
	return loadResumeSpec(spec, res)
}

// Name of the torrent.
// For magnet downloads name can change after metadata is downloaded but this method still returns the initial name.
func (t *Torrent) Name() string {
	return t.name
}

// InfoHash string encoded in hex.
// InfoHash is a unique value that identifies the files in torrent.
func (t *Torrent) InfoHash() string {
	return hex.EncodeToString(t.infoHash[:])
}

// SetResume adds resume capability to the torrent.
// It must be called before Start() is called.
func (t *Torrent) SetResume(res resume.DB) error {
	spec, err := res.Read()
	if err != nil {
		return err
	}
	if spec == nil {
		return t.writeResume(res)
	}
	t2, err := loadResumeSpec(spec, res)
	if err != nil {
		return err
	}
	if t.InfoHash() != t2.InfoHash() {
		t2.Close()
		return errors.New("invalid resume file (info hashes does not match)")
	}
	t.Close()
	*t = *t2
	return nil
}

func loadResumeSpec(spec *resume.Spec, res resume.DB) (*Torrent, error) {
	var err error
	dspec := &downloadSpec{
		port:     spec.Port,
		trackers: spec.Trackers,
		name:     spec.Name,
		resume:   res,
	}
	copy(dspec.infoHash[:], spec.InfoHash)
	if len(spec.Info) > 0 {
		dspec.info, err = metainfo.NewInfo(spec.Info)
		if err != nil {
			return nil, err
		}
		if len(spec.Bitfield) > 0 {
			dspec.bitfield = bitfield.New(dspec.info.NumPieces)
			copy(dspec.bitfield.Bytes(), spec.Bitfield)
		}
	}
	switch spec.StorageType {
	case filestorage.StorageType:
		dspec.storage = &filestorage.FileStorage{}
	default:
		return nil, errors.New("unknown storage type: " + spec.StorageType)
	}
	err = dspec.storage.Load(spec.StorageArgs)
	if err != nil {
		return nil, err
	}
	return newTorrent(dspec)
}

func (t *Torrent) writeResume(res resume.DB) error {
	rspec := &resume.Spec{
		InfoHash:    t.infoHash[:],
		Port:        t.port,
		Name:        t.name,
		Trackers:    t.trackers,
		StorageType: t.storage.Type(),
		StorageArgs: t.storage.Args(),
	}
	if t.info != nil {
		rspec.Info = t.info.Bytes
	}
	if t.bitfield != nil {
		rspec.Bitfield = t.bitfield.Bytes()
	}
	return res.Write(rspec)
}

func newTorrent(spec *downloadSpec) (*Torrent, error) {
	logName := spec.name
	if len(logName) > 8 {
		logName = logName[:8]
	}
	d := &Torrent{
		downloadSpec:              spec,
		log:                       logger.New("download " + logName),
		peerDisconnectedC:         make(chan *peer.Peer),
		messages:                  make(chan peer.Message),
		peers:                     make(map[*peer.Peer]struct{}),
		incomingPeers:             make(map[*peer.Peer]struct{}),
		outgoingPeers:             make(map[*peer.Peer]struct{}),
		peersSnubbed:              make(map[*peer.Peer]struct{}),
		pieceDownloaders:          make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		pieceDownloadersSnubbed:   make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		pieceDownloadersChoked:    make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		snubbedPieceDownloaderC:   make(chan *piecedownloader.PieceDownloader),
		snubbedInfoDownloaderC:    make(chan *infodownloader.InfoDownloader),
		infoDownloaders:           make(map[*peer.Peer]*infodownloader.InfoDownloader),
		infoDownloadersSnubbed:    make(map[*peer.Peer]*infodownloader.InfoDownloader),
		completeC:                 make(chan struct{}),
		closeC:                    make(chan struct{}),
		doneC:                     make(chan struct{}),
		startCommandC:             make(chan struct{}),
		stopCommandC:              make(chan struct{}),
		statsCommandC:             make(chan statsRequest),
		notifyErrorCommandC:       make(chan notifyErrorCommand),
		addrsFromTrackers:         make(chan []*net.TCPAddr),
		addrList:                  addrlist.New(),
		peerIDs:                   make(map[[20]byte]struct{}),
		incomingConnC:             make(chan net.Conn),
		sKeyHash:                  mse.HashSKey(spec.infoHash[:]),
		infoDownloaderDoneC:       make(chan *infodownloader.InfoDownloader),
		pieceDownloaderDoneC:      make(chan *piecedownloader.PieceDownloader),
		incomingHandshakers:       make(map[string]*incominghandshaker.IncomingHandshaker),
		outgoingHandshakers:       make(map[string]*outgoinghandshaker.OutgoingHandshaker),
		incomingHandshakerResultC: make(chan incominghandshaker.Result),
		outgoingHandshakerResultC: make(chan outgoinghandshaker.Result),
		announcerRequestC:         make(chan *announcer.Request),
		allocatorProgressC:        make(chan allocator.Progress),
		allocatorResultC:          make(chan allocator.Result),
		verifierProgressC:         make(chan verifier.Progress),
		verifierResultC:           make(chan verifier.Result),
	}
	copy(d.peerID[:], peerIDPrefix)
	_, err := rand.Read(d.peerID[len(peerIDPrefix):]) // nolint: gosec
	if err != nil {
		return nil, err
	}
	d.trackersInstances, err = parseTrackers(d.trackers, d.log)
	if err != nil {
		return nil, err
	}
	go d.run()
	return d, nil
}

func parseTrackers(trackers []string, log logger.Logger) ([]tracker.Tracker, error) {
	var ret []tracker.Tracker
	for _, s := range trackers {
		u, err := url.Parse(s)
		if err != nil {
			log.Warningln("cannot parse tracker url:", err)
			continue
		}
		switch u.Scheme {
		case "http", "https":
			ret = append(ret, httptracker.New(u))
		case "udp":
			ret = append(ret, udptracker.New(u))
		default:
			log.Warningln("unsupported tracker scheme: %s", u.Scheme)
		}
	}
	if len(ret) == 0 {
		return nil, errors.New("no tracker found")
	}
	return ret, nil
}
<commit_msg>tuning<commit_after>// Package torrent provides a torrent client implementation.
package torrent

import (
	"encoding/hex"
	"errors"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"

	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/resume"
	"github.com/cenkalti/rain/storage"
	"github.com/cenkalti/rain/storage/filestorage"
	"github.com/cenkalti/rain/torrent/internal/acceptor"
	"github.com/cenkalti/rain/torrent/internal/addrlist"
	"github.com/cenkalti/rain/torrent/internal/allocator"
	"github.com/cenkalti/rain/torrent/internal/announcer"
	"github.com/cenkalti/rain/torrent/internal/bitfield"
	"github.com/cenkalti/rain/torrent/internal/handshaker/incominghandshaker"
	"github.com/cenkalti/rain/torrent/internal/handshaker/outgoinghandshaker"
	"github.com/cenkalti/rain/torrent/internal/infodownloader"
	"github.com/cenkalti/rain/torrent/internal/magnet"
	"github.com/cenkalti/rain/torrent/internal/metainfo"
	"github.com/cenkalti/rain/torrent/internal/mse"
	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/piece"
	"github.com/cenkalti/rain/torrent/internal/piecedownloader"
	"github.com/cenkalti/rain/torrent/internal/torrentdata"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/httptracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/udptracker"
	"github.com/cenkalti/rain/torrent/internal/verifier"
)

const (
	maxPeerDial            = 25
	maxPeerAccept          = 25
	parallelPieceDownloads = 50 // running downloads, snubbed and choked peers don't count
	parallelInfoDownloads  = 1  // running downloads, snubbed peers don't count
)

var (
	// http://www.bittorrent.org/beps/bep_0020.html
	peerIDPrefix = []byte("-RN" + clientversion.Version + "-")
)

// Torrent connects to peers and downloads files from swarm.
type Torrent struct {
	*downloadSpec

	// Unique peer ID is generated per downloader.
	peerID [20]byte

	// Data provides IO access to pieces in torrent.
	data *torrentdata.Data

	// Contains state about the pieces in torrent.
	pieces []*piece.Piece

	// Contains pieces in sorted order for piece selection function.
	sortedPieces []*piece.Piece

	// Peers are sent to this channel when they are disconnected.
	peerDisconnectedC chan *peer.Peer

	// All messages coming from peers are sent to this channel.
	messages chan peer.Message

	// We keep connected peers in this map after they complete handshake phase.
	peers         map[*peer.Peer]struct{}
	incomingPeers map[*peer.Peer]struct{}
	outgoingPeers map[*peer.Peer]struct{}
	peersSnubbed  map[*peer.Peer]struct{}

	// Active piece downloads are kept in this map.
	pieceDownloaders        map[*peer.Peer]*piecedownloader.PieceDownloader
	pieceDownloadersSnubbed map[*peer.Peer]*piecedownloader.PieceDownloader
	pieceDownloadersChoked  map[*peer.Peer]*piecedownloader.PieceDownloader

	// When piece downloaders detects that a peer has snubbed us, it will send a signal to this channel.
	snubbedPieceDownloaderC chan *piecedownloader.PieceDownloader
	snubbedInfoDownloaderC  chan *infodownloader.InfoDownloader

	// Active metadata downloads are kept in this map.
	infoDownloaders        map[*peer.Peer]*infodownloader.InfoDownloader
	infoDownloadersSnubbed map[*peer.Peer]*infodownloader.InfoDownloader

	// A peer is optimistically unchoked regardless of their download rate.
	optimisticUnchokedPeer *peer.Peer

	// This channel is closed once all pieces are downloaded and verified.
	completeC chan struct{}

	// If any unrecoverable error occurs, it will be sent to this channel and download will be stopped.
	errC chan error

	// When Stop() is called, it will close this channel to signal run() function to stop.
	closeC chan struct{}

	// This channel will be closed after run loop exists.
	doneC chan struct{}

	// These are the channels for sending a message to run() loop.
	statsCommandC       chan statsRequest       // Stats()
	startCommandC       chan struct{}           // Start()
	stopCommandC        chan struct{}           // Stop()
	notifyErrorCommandC chan notifyErrorCommand // NotifyError()

	// Trackers send announce responses to this channel.
	addrsFromTrackers chan []*net.TCPAddr

	// Keeps a list of peer addresses to connect.
	addrList *addrlist.AddrList

	// New raw connections created by OutgoingHandshaker are sent to here.
	incomingConnC chan net.Conn

	// Keep a set of peer IDs to block duplicate connections.
	peerIDs map[[20]byte]struct{}

	// Listens for incoming peer connections.
	acceptor *acceptor.Acceptor

	// Special hash of info hash for encypted connection handshake.
	sKeyHash [20]byte

	// Tracker implementations for giving to announcers.
	trackersInstances []tracker.Tracker

	// Announces the status of torrent to trackers to get peer addresses.
	announcers []*announcer.Announcer

	// List of peers in handshake state.
	incomingHandshakers map[string]*incominghandshaker.IncomingHandshaker
	outgoingHandshakers map[string]*outgoinghandshaker.OutgoingHandshaker

	// Handshake results are sent to these channels by handshakers.
	incomingHandshakerResultC chan incominghandshaker.Result
	outgoingHandshakerResultC chan outgoinghandshaker.Result

	// When metadata of the torrent downloaded completely, a message is sent to this channel.
	infoDownloaderDoneC chan *infodownloader.InfoDownloader

	// When a piece is downloaded completely a message is sent to this channel.
	pieceDownloaderDoneC chan *piecedownloader.PieceDownloader

	// Announcers send a request to this channel to get information about the torrent.
	announcerRequestC chan *announcer.Request

	// A timer that ticks periodically to keep a certain number of peers unchoked.
	unchokeTimer  *time.Ticker
	unchokeTimerC <-chan time.Time

	// A timer that ticks periodically to keep a random peer unchoked regardless of its upload rate.
	optimisticUnchokeTimer  *time.Ticker
	optimisticUnchokeTimerC <-chan time.Time

	allocator          *allocator.Allocator
	allocatorProgressC chan allocator.Progress
	allocatorResultC   chan allocator.Result

	verifier          *verifier.Verifier
	verifierProgressC chan verifier.Progress
	verifierResultC   chan verifier.Result

	bytesDownloaded int64
	bytesWasted     int64

	log logger.Logger
}

// New returns a new torrent by reading a torrent metainfo file.
func New(r io.Reader, port int, sto storage.Storage) (*Torrent, error) {
	m, err := metainfo.New(r)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.Info.Hash,
		trackers: m.GetTrackers(),
		name:     m.Info.Name,
		port:     port,
		storage:  sto,
		info:     m.Info,
	}
	return newTorrent(spec)
}

// NewMagnet returns a new torrent by parsing a magnet link.
func NewMagnet(link string, port int, sto storage.Storage) (*Torrent, error) {
	m, err := magnet.New(link)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.InfoHash,
		trackers: m.Trackers,
		name:     m.Name,
		port:     port,
		storage:  sto,
	}
	return newTorrent(spec)
}

// NewResume returns a new torrent by loading all info from a resume.DB.
func NewResume(res resume.DB) (*Torrent, error) {
	spec, err := res.Read()
	if err != nil {
		return nil, err
	}
	if spec == nil {
		return nil, errors.New("no resume info")
	}
	return loadResumeSpec(spec, res)
}

// Name of the torrent.
// For magnet downloads name can change after metadata is downloaded but this method still returns the initial name.
func (t *Torrent) Name() string {
	return t.name
}

// InfoHash string encoded in hex.
// InfoHash is a unique value that identifies the files in torrent.
func (t *Torrent) InfoHash() string {
	return hex.EncodeToString(t.infoHash[:])
}

// SetResume adds resume capability to the torrent.
// It must be called before Start() is called.
func (t *Torrent) SetResume(res resume.DB) error {
	spec, err := res.Read()
	if err != nil {
		return err
	}
	if spec == nil {
		return t.writeResume(res)
	}
	t2, err := loadResumeSpec(spec, res)
	if err != nil {
		return err
	}
	if t.InfoHash() != t2.InfoHash() {
		t2.Close()
		return errors.New("invalid resume file (info hashes does not match)")
	}
	t.Close()
	*t = *t2
	return nil
}

func loadResumeSpec(spec *resume.Spec, res resume.DB) (*Torrent, error) {
	var err error
	dspec := &downloadSpec{
		port:     spec.Port,
		trackers: spec.Trackers,
		name:     spec.Name,
		resume:   res,
	}
	copy(dspec.infoHash[:], spec.InfoHash)
	if len(spec.Info) > 0 {
		dspec.info, err = metainfo.NewInfo(spec.Info)
		if err != nil {
			return nil, err
		}
		if len(spec.Bitfield) > 0 {
			dspec.bitfield = bitfield.New(dspec.info.NumPieces)
			copy(dspec.bitfield.Bytes(), spec.Bitfield)
		}
	}
	switch spec.StorageType {
	case filestorage.StorageType:
		dspec.storage = &filestorage.FileStorage{}
	default:
		return nil, errors.New("unknown storage type: " + spec.StorageType)
	}
	err = dspec.storage.Load(spec.StorageArgs)
	if err != nil {
		return nil, err
	}
	return newTorrent(dspec)
}

func (t *Torrent) writeResume(res resume.DB) error {
	rspec := &resume.Spec{
		InfoHash:    t.infoHash[:],
		Port:        t.port,
		Name:        t.name,
		Trackers:    t.trackers,
		StorageType: t.storage.Type(),
		StorageArgs: t.storage.Args(),
	}
	if t.info != nil {
		rspec.Info = t.info.Bytes
	}
	if t.bitfield != nil {
		rspec.Bitfield = t.bitfield.Bytes()
	}
	return res.Write(rspec)
}

func newTorrent(spec *downloadSpec) (*Torrent, error) {
	logName := spec.name
	if len(logName) > 8 {
		logName = logName[:8]
	}
	d := &Torrent{
		downloadSpec:              spec,
		log:                       logger.New("download " + logName),
		peerDisconnectedC:         make(chan *peer.Peer),
		messages:                  make(chan peer.Message),
		peers:                     make(map[*peer.Peer]struct{}),
		incomingPeers:             make(map[*peer.Peer]struct{}),
		outgoingPeers:             make(map[*peer.Peer]struct{}),
		peersSnubbed:              make(map[*peer.Peer]struct{}),
		pieceDownloaders:          make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		pieceDownloadersSnubbed:   make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		pieceDownloadersChoked:    make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		snubbedPieceDownloaderC:   make(chan *piecedownloader.PieceDownloader),
		snubbedInfoDownloaderC:    make(chan *infodownloader.InfoDownloader),
		infoDownloaders:           make(map[*peer.Peer]*infodownloader.InfoDownloader),
		infoDownloadersSnubbed:    make(map[*peer.Peer]*infodownloader.InfoDownloader),
		completeC:                 make(chan struct{}),
		closeC:                    make(chan struct{}),
		doneC:                     make(chan struct{}),
		startCommandC:             make(chan struct{}),
		stopCommandC:              make(chan struct{}),
		statsCommandC:             make(chan statsRequest),
		notifyErrorCommandC:       make(chan notifyErrorCommand),
		addrsFromTrackers:         make(chan []*net.TCPAddr),
		addrList:                  addrlist.New(),
		peerIDs:                   make(map[[20]byte]struct{}),
		incomingConnC:             make(chan net.Conn),
		sKeyHash:                  mse.HashSKey(spec.infoHash[:]),
		infoDownloaderDoneC:       make(chan *infodownloader.InfoDownloader),
		pieceDownloaderDoneC:      make(chan *piecedownloader.PieceDownloader),
		incomingHandshakers:       make(map[string]*incominghandshaker.IncomingHandshaker),
		outgoingHandshakers:       make(map[string]*outgoinghandshaker.OutgoingHandshaker),
		incomingHandshakerResultC: make(chan incominghandshaker.Result),
		outgoingHandshakerResultC: make(chan outgoinghandshaker.Result),
		announcerRequestC:         make(chan *announcer.Request),
		allocatorProgressC:        make(chan allocator.Progress),
		allocatorResultC:          make(chan allocator.Result),
		verifierProgressC:         make(chan verifier.Progress),
		verifierResultC:           make(chan verifier.Result),
	}
	copy(d.peerID[:], peerIDPrefix)
	_, err := rand.Read(d.peerID[len(peerIDPrefix):]) // nolint: gosec
	if err != nil {
		return nil, err
	}
	d.trackersInstances, err = parseTrackers(d.trackers, d.log)
	if err != nil {
		return nil, err
	}
	go d.run()
	return d, nil
}

func parseTrackers(trackers []string, log logger.Logger) ([]tracker.Tracker, error) {
	var ret []tracker.Tracker
	for _, s := range trackers {
		u, err := url.Parse(s)
		if err != nil {
			log.Warningln("cannot parse tracker url:", err)
			continue
		}
		switch u.Scheme {
		case "http", "https":
			ret = append(ret, httptracker.New(u))
		case "udp":
			ret = append(ret, udptracker.New(u))
		default:
			log.Warningln("unsupported tracker scheme: %s", u.Scheme)
		}
	}
	if len(ret) == 0 {
		return nil, errors.New("no tracker found")
	}
	return ret, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"fmt"
	"log"
	"os"
)

type kml struct {
	Document struct {
		Placemarks []placemark `xml:"Placemark"`
	}
}

type placemark struct {
	Name  string `xml:"name"`
	Desc  string `xml:"description"`
	Point point  `xml:"Point"`
}

type point struct {
	Coord coord `xml:"coordinates"`
}

type coord string

const (
	communityClubsKML = "communityclubs.kml"
)

func (k *kml) debugPrint() {
	for _, p := range k.Document.Placemarks {
		fmt.Println(p.Name)
		fmt.Println(p.Point)
		fmt.Println(p.Desc)
	}
}

func main() {
	f, err := os.Open(communityClubsKML)
	if err != nil {
		log.Fatal(err)
	}
	k := kml{}
	xml.NewDecoder(f).Decode(&k)
	k.debugPrint()
}
<commit_msg>Finish up KML parsing<commit_after>package main

import (
	"encoding/xml"
	"fmt"
	"log"
	"os"
	"regexp"
)

type kml struct {
	Document struct {
		Placemarks []placemark `xml:"Placemark"`
	}
}

type placemark struct {
	Name  string `xml:"name"`
	Desc  desc   `xml:"description"`
	Point point  `xml:"Point"`
}

type desc struct {
	Lines []struct {
		Colour struct {
			Bold string `xml:"b"`
		} `xml:"font"`
	} `xml:"p"`
}

type point struct {
	Coord string `xml:"coordinates"`
}

const (
	communityClubsKML = "communityclubs.kml"
)

func (k *kml) debugPrint() {
	for _, p := range k.Document.Placemarks {
		fmt.Println(p.Name)
		fmt.Println(p.Point.Coord)
		for _, l := range p.Desc.Lines {
			fmt.Println(l.Colour.Bold)
		}
	}
}

type CommunityClub struct {
	Name  string `json:"name"`
	Point string `json:"point"`
	Desc  string `json:"desc"`
}

var (
	descMatch = regexp.MustCompile(`Description - (.*)`)
)

func NewCommunityClubs(k *kml) []CommunityClub {
	var clubs []CommunityClub
	for _, p := range k.Document.Placemarks {
		c := CommunityClub{
			Name:  p.Name,
			Point: p.Point.Coord,
		}
		for _, l := range p.Desc.Lines {
			s := l.Colour.Bold
			if m := descMatch.FindStringSubmatch(s); m != nil {
				fmt.Println(m)
				c.Desc = m[1]
			}
		}
		clubs = append(clubs, c)
	}
	return clubs
}

func printCommunityClubs(clubs []CommunityClub) {
	for _, c := range clubs {
		fmt.Println(c.Name)
		fmt.Println(c.Desc)
		fmt.Println(c.Point)
	}
}

func main() {
	f, err := os.Open(communityClubsKML)
	if err != nil {
		log.Fatal(err)
	}
	k := kml{}
	xml.NewDecoder(f).Decode(&k)
	clubs := NewCommunityClubs(&k)
	printCommunityClubs(clubs)
}
<|endoftext|>


<commit_before>// git.go
package translation

import (
	"bitbucket.org/kardianos/rsync"
	"github.com/libgit2/git2go"
	"github.com/rgeorgiev583/gonflator/remote"
)

const (
	chanCap  = 1000
	sliceCap = 1000
)

type SyntheticRemote struct {
	URL      url.URL
	Protocol remote.Protocol
}

type SyntheticRemoteCollection map[string]*SyntheticRemote

type GitRepository struct {
	git2go.Repository
	Tree             ConfigurationTree
	SyntheticRemotes SyntheticRemoteCollection
}

func (gr *GitRepository) GetRdiff(diff chan<- git2go.DiffDelta) (rdiff <-chan Delta, err error) {
	rdiff = make(chan Delta, chanCap)

	go func() {
		defer func() {
			close(rdiff)
		}()

		for delta := range diff {
			switch delta.Status {
			case git2go.DeltaUnmodified:
				rdiff <- &Delta{
					NewPath: delta.NewFile.Path,
					Type:    DeltaUnmodified,
				}
			case git2go.DeltaAdded:
				blob, err := gr.LookupBlob(delta.NewFile.Oid)
				if err != nil {
					return
				}

				rdiff <- &Delta{
					Operation: &rsync.Operation{Data: blob.Contents()},
					NewPath:   delta.NewFile.Path,
					Type:      DeltaAdded,
				}
			case git2go.DeltaDeleted:
				rdiff <- &Delta{
					OldPath: delta.OldFile.Path,
					Type:    DeltaDeleted,
				}
			case git2go.DeltaModified:
				newBlob, err := gr.LookupBlob(delta.NewFile.Oid)
				if err != nil {
					return
				}

				oldBlob, err := gr.LookupBlob(delta.OldFile.Oid)
				if err != nil {
					return
				}

				rdiffMaker := &rsync.Rsync{}
				oldReader := bytes.NewReader(oldBlob.Contents())
				newReader := bytes.NewReader(newBlob.Contents())
				signature := new([]BlockHash, 0, sliceCapacity)
				err = rsync.CreateSignature(oldReader, func(bh BlockHash) error {
					append(signature, bh)
					return nil
				})
				if err != nil {
					return
				}

				err = rsync.CreateDelta(newReader, signature, func(op Operation) error {
					rdiff <- &Delta{
						Operation: op,
						OldPath:   delta.OldFile.Path,
						NewPath:   delta.NewFile.Path,
						Type:      DeltaModified,
					}
					return nil
				})
			case git2go.DeltaRenamed:
				rdiff <- &Delta{
					OldPath: delta.OldFile.Path,
					NewPath: delta.NewFile.Path,
					Type:    DeltaRenamed,
				}
			case git2go.DeltaCopied:
				rdiff <- &Delta{
					OldPath: delta.OldFile.Path,
					Path:    delta.NewFile.Path,
					Type:    DeltaCopied,
				}
			}
		}
	}()

	return
}

func (gr *GitRepository) DiffReference(ref *git2go.Reference) (*git2go.Diff, error) {
	return gr.DiffTreeToWorkdirWithIndex(ref.Target(), nil)
}

func (gr *GitRepository) DiffRemote(ref *git2go.Reference, remoteName string) (translatedDiff <-chan Delta, err error) {
	gitDiff, err := gr.DiffReference(ref)
	if err != nil {
		return
	}

	diff := make(chan git2go.DiffDelta, chanCap)
	defer close(diff)
	deltaCollector := func(delta git2go.DiffDelta, _ float64) (git2go.DiffForEachHunkCallback, error) {
		diff <- delta
		return
	}
	err = diff.ForEach(deltaCollector, git2go.DiffDetailFiles)
	if err != nil {
		return
	}

	remote, err := gr.SyntheticRemotes[remoteName]
	if err != nil {
		return
	}

	comm := remote.DialURL(remote.URL, remote.Protocol)
	return gr.GetRdiff(rdiff)
}

func (gr *GitRepository) Add(name string, rawurl string, proto remote.Protocol) (err error) {
	url, err := url.Parse(rawurl)
	if err != nil {
		return
	}

	gr.SyntheticRemotes[name] = &SyntheticRemote{
		URL:      url,
		Protocol: proto,
	}
}
<commit_msg>optimize: reduced a goroutine to a go-statement<commit_after>// git.go
package translation

import (
	"bitbucket.org/kardianos/rsync"
	"github.com/libgit2/git2go"
	"github.com/rgeorgiev583/gonflator/remote"
)

const (
	chanCap  = 1000
	sliceCap = 1000
)

type SyntheticRemote struct {
	URL      url.URL
	Protocol remote.Protocol
}

type SyntheticRemoteCollection map[string]*SyntheticRemote

type GitRepository struct {
	git2go.Repository
	Tree             ConfigurationTree
	SyntheticRemotes SyntheticRemoteCollection
}

func (gr *GitRepository) GetRdiff(diff chan<- git2go.DiffDelta) (rdiff <-chan Delta, err error) {
	rdiff = make(chan Delta, chanCap)

	go func() {
		defer close(rdiff)

		for delta := range diff {
			switch delta.Status {
			case git2go.DeltaUnmodified:
				rdiff <- &Delta{
					NewPath: delta.NewFile.Path,
					Type:    DeltaUnmodified,
				}
			case git2go.DeltaAdded:
				blob, err := gr.LookupBlob(delta.NewFile.Oid)
				if err != nil {
					return
				}

				rdiff <- &Delta{
					Operation: &rsync.Operation{Data: blob.Contents()},
					NewPath:   delta.NewFile.Path,
					Type:      DeltaAdded,
				}
			case git2go.DeltaDeleted:
				rdiff <- &Delta{
					OldPath: delta.OldFile.Path,
					Type:    DeltaDeleted,
				}
			case git2go.DeltaModified:
				newBlob, err := gr.LookupBlob(delta.NewFile.Oid)
				if err != nil {
					return
				}

				oldBlob, err := gr.LookupBlob(delta.OldFile.Oid)
				if err != nil {
					return
				}

				rdiffMaker := &rsync.Rsync{}
				oldReader := bytes.NewReader(oldBlob.Contents())
				newReader := bytes.NewReader(newBlob.Contents())
				signature := new([]BlockHash, 0, sliceCapacity)
				err = rsync.CreateSignature(oldReader, func(bh BlockHash) error {
					append(signature, bh)
					return nil
				})
				if err != nil {
					return
				}

				err = rsync.CreateDelta(newReader, signature, func(op Operation) error {
					rdiff <- &Delta{
						Operation: op,
						OldPath:   delta.OldFile.Path,
						NewPath:   delta.NewFile.Path,
						Type:      DeltaModified,
					}
					return nil
				})
			case git2go.DeltaRenamed:
				rdiff <- &Delta{
					OldPath: delta.OldFile.Path,
					NewPath: delta.NewFile.Path,
					Type:    DeltaRenamed,
				}
			case git2go.DeltaCopied:
				rdiff <- &Delta{
					OldPath: delta.OldFile.Path,
					Path:    delta.NewFile.Path,
					Type:    DeltaCopied,
				}
			}
		}
	}()

	return
}

func (gr *GitRepository) DiffReference(ref *git2go.Reference) (*git2go.Diff, error) {
	return gr.DiffTreeToWorkdirWithIndex(ref.Target(), nil)
}

func (gr *GitRepository) DiffRemote(ref *git2go.Reference, remoteName string) (translatedDiff <-chan Delta, err error) {
	gitDiff, err := gr.DiffReference(ref)
	if err != nil {
		return
	}

	diff := make(chan git2go.DiffDelta, chanCap)
	defer close(diff)
	deltaCollector := func(delta git2go.DiffDelta, _ float64) (git2go.DiffForEachHunkCallback, error) {
		diff <- delta
		return
	}
	err = diff.ForEach(deltaCollector, git2go.DiffDetailFiles)
	if err != nil {
		return
	}

	remote, err := gr.SyntheticRemotes[remoteName]
	if err != nil {
		return
	}

	comm := remote.DialURL(remote.URL, remote.Protocol)
	return gr.GetRdiff(rdiff)
}

func (gr *GitRepository) Add(name string, rawurl string, proto remote.Protocol) (err error) {
	url, err := url.Parse(rawurl)
	if err != nil {
		return
	}

	gr.SyntheticRemotes[name] = &SyntheticRemote{
		URL:      url,
		Protocol: proto,
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Mangos Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use 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 ws implements a simple WebSocket transport for mangos.
// This transport is considered EXPERIMENTAL.
package ws

import (
	"crypto/tls"
	"github.com/gorilla/websocket"
	"net"
	"net/http"
	"net/url"
	"strings"

	"github.com/go-mangos/mangos"
	"sync"
)

// Some special options
const (
	// OptionWebSocketMux is a retrieve-only property used to obtain
	// the *http.ServeMux instance associated with the server.  This
	// can be used to subsequently register additional handlers for
	// different URIs.  This option is only valid on a Listener.
	// Generally you use this option when you want to use the standard
	// mangos Listen() method to start up the server.
	OptionWebSocketMux = "WEBSOCKET-MUX"

	// OptionWebSocketHandler is used to obtain the underlying
	// http.Handler (websocket.Server) object, so you can use this
	// on your own http.Server instances.  It is a gross error to use
	// the value returned by this method on an http server if the
	// server is also started with mangos Listen().  This means that you
	// will use at most either this option, or OptionWebSocketMux, but
	// never both.  This option is only valid on a listener.
	OptionWebSocketHandler = "WEBSOCKET-HANDLER"
)

type options map[string]interface{}

// GetOption retrieves an option value.
func (o options) get(name string) (interface{}, error) {
	if o == nil {
		return nil, mangos.ErrBadOption
	}
	v, ok := o[name]
	if !ok {
		return nil, mangos.ErrBadOption
	}
	return v, nil
}

// SetOption sets an option.  We have none, so just ErrBadOption.
func (o options) set(name string, val interface{}) error {
	switch name {
	case mangos.OptionNoDelay:
		fallthrough
	case mangos.OptionKeepAlive:
		switch v := val.(type) {
		case bool:
			o[name] = v
			return nil
		default:
			return mangos.ErrBadValue
		}
	case mangos.OptionTLSConfig:
		switch v := val.(type) {
		case *tls.Config:
			o[name] = v
			return nil
		default:
			return mangos.ErrBadValue
		}
	}
	return mangos.ErrBadOption
}

// wsPipe implements the Pipe interface on a websocket
type wsPipe struct {
	ws    *websocket.Conn
	proto mangos.Protocol
	addr  string
	open  bool
	wg    sync.WaitGroup
	props map[string]interface{}
	iswss bool
	dtype int
}

type wsTran int

func (w *wsPipe) Recv() (*mangos.Message, error) {

	// We ignore the message type for receive.
	_, body, err := w.ws.ReadMessage()
	if err != nil {
		return nil, err
	}
	msg := mangos.NewMessage(0)
	msg.Body = body
	return msg, nil
}

func (w *wsPipe) Send(m *mangos.Message) error {

	var buf []byte

	if m.Expired() {
		m.Free()
		return nil
	}
	if len(m.Header) > 0 {
		buf = make([]byte, 0, len(m.Header)+len(m.Body))
		buf = append(buf, m.Header...)
		buf = append(buf, m.Body...)
	} else {
		buf = m.Body
	}
	if err := w.ws.WriteMessage(w.dtype, buf); err != nil {
		return err
	}
	m.Free()
	return nil
}

func (w *wsPipe) LocalProtocol() uint16 {
	return w.proto.Number()
}

func (w *wsPipe) RemoteProtocol() uint16 {
	return w.proto.PeerNumber()
}

func (w *wsPipe) Close() error {
	w.open = false
	w.ws.Close()
	w.wg.Done()
	return nil
}

func (w *wsPipe) IsOpen() bool {
	return w.open
}

func (w *wsPipe) GetProp(name string) (interface{}, error) {
	if v, ok := w.props[name]; ok {
		return v, nil
	}
	return nil, mangos.ErrBadProperty
}

type dialer struct {
	addr  string // url
	proto mangos.Protocol
	opts  options
	iswss bool
	maxrx int
}

func (d *dialer) Dial() (mangos.Pipe, error) {
	var w *wsPipe

	wd := &websocket.Dialer{}

	wd.Subprotocols = []string{d.proto.PeerName() + ".sp.nanomsg.org"}
	if v, ok := d.opts[mangos.OptionTLSConfig]; ok {
		wd.TLSClientConfig = v.(*tls.Config)
	}

	w = &wsPipe{proto: d.proto, addr: d.addr, open: true}
	w.dtype = websocket.BinaryMessage
	w.props = make(map[string]interface{})

	var err error
	if w.ws, _, err = wd.Dial(d.addr, nil); err != nil {
		return nil, err
	}
	w.ws.SetReadLimit(int64(d.maxrx))
	w.props[mangos.PropLocalAddr] = w.ws.LocalAddr()
	w.props[mangos.PropRemoteAddr] = w.ws.RemoteAddr()

	w.wg.Add(1)
	return w, nil
}

func (d *dialer) SetOption(n string, v interface{}) error {
	return d.opts.set(n, v)
}

func (d *dialer) GetOption(n string) (interface{}, error) {
	return d.opts.get(n)
}

type listener struct {
	pending  []*wsPipe
	lock     sync.Mutex
	cv       sync.Cond
	running  bool
	noserve  bool
	addr     string
	ug       websocket.Upgrader
	htsvr    *http.Server
	mux      *http.ServeMux
	url      *url.URL
	listener net.Listener
	proto    mangos.Protocol
	opts     options
	iswss    bool
	maxrx    int
}

func (l *listener) SetOption(n string, v interface{}) error {
	return l.opts.set(n, v)
}

func (l *listener) GetOption(n string) (interface{}, error) {
	switch n {
	case OptionWebSocketMux:
		return l.mux, nil
	case OptionWebSocketHandler:
		// Caller intends to use use in his own server, so mark
		// us running.  If he didn't mean this, the side effect is
		// that Accept() will appear to hang, even though Listen()
		// is not called yet.
		l.running = true
		l.noserve = true
		return l, nil
	}
	return l.opts.get(n)
}

func (l *listener) Listen() error {
	var taddr *net.TCPAddr
	var err error
	var tcfg *tls.Config

	if l.noserve {
		// The HTTP framework is going to call us, so we use that rather than
		// listening on our own.  We just fake this out.
		return nil
	}
	if l.iswss {
		v, ok := l.opts[mangos.OptionTLSConfig]
		if !ok || v == nil {
			return mangos.ErrTLSNoConfig
		}
		tcfg = v.(*tls.Config)
		if tcfg.Certificates == nil || len(tcfg.Certificates) == 0 {
			return mangos.ErrTLSNoCert
		}
	}

	// We listen separately, that way we can catch and deal with the
	// case of a port already in use.  This also lets us configure
	// properties of the underlying TCP connection.

	if taddr, err = mangos.ResolveTCPAddr(l.url.Host); err != nil {
		return err
	}

	if tlist, err := net.ListenTCP("tcp", taddr); err != nil {
		return err
	} else if l.iswss {
		l.listener = tls.NewListener(tlist, tcfg)
	} else {
		l.listener = tlist
	}
	l.pending = nil
	l.running = true

	l.htsvr = &http.Server{Addr: l.url.Host, Handler: l.mux}

	go l.htsvr.Serve(l.listener)

	return nil
}

func (l *listener) Accept() (mangos.Pipe, error) {
	var w *wsPipe

	l.lock.Lock()
	defer l.lock.Unlock()

	for {
		if !l.running {
			return nil, mangos.ErrClosed
		}
		if len(l.pending) == 0 {
			l.cv.Wait()
			continue
		}
		w = l.pending[len(l.pending)-1]
		l.pending = l.pending[:len(l.pending)-1]
		break
	}

	return w, nil
}

func (l *listener) handler(ws *websocket.Conn, req *http.Request) {
	l.lock.Lock()

	if !l.running {
		ws.Close()
		l.lock.Unlock()
		return
	}

	if ws.Subprotocol() != l.proto.Name()+".sp.nanomsg.org" {
		ws.Close()
		l.lock.Unlock()
		return
	}

	w := &wsPipe{ws: ws, addr: l.addr, proto: l.proto, open: true}
	w.dtype = websocket.BinaryMessage
	w.iswss = l.iswss
	w.ws.SetReadLimit(int64(l.maxrx))

	w.props = make(map[string]interface{})
	w.props[mangos.PropLocalAddr] = ws.LocalAddr()
	w.props[mangos.PropRemoteAddr] = ws.RemoteAddr()

	if req.TLS != nil {
		w.props[mangos.PropTLSConnState] = *req.TLS
	}

	w.wg.Add(1)
	l.pending = append(l.pending, w)
	l.cv.Broadcast()
	l.lock.Unlock()

	// We must not return before the socket is closed, because
	// our caller will close the websocket on our return.
	w.wg.Wait()
}

func (l *listener) Handle(pattern string, handler http.Handler) {
	l.mux.Handle(pattern, handler)
}

func (l *listener) HandleFunc(pattern string, handler http.HandlerFunc) {
	l.mux.HandleFunc(pattern, handler)
}

func (l *listener) Close() error {
	l.lock.Lock()
	defer l.lock.Unlock()
	if !l.running {
		return mangos.ErrClosed
	}
	if l.listener != nil {
		l.listener.Close()
	}
	l.running = false
	l.cv.Broadcast()
	for _, ws := range l.pending {
		ws.Close()
	}
	l.pending = nil
	return nil
}

func (l *listener) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ws, err := l.ug.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	l.handler(ws, r)
}

func (l *listener) Address() string {
	return l.url.String()
}

func (wsTran) Scheme() string {
	return "ws"
}

func (wsTran) NewDialer(addr string, sock mangos.Socket) (mangos.PipeDialer, error) {
	iswss := strings.HasPrefix(addr, "wss://")
	opts := make(map[string]interface{})

	opts[mangos.OptionNoDelay] = true
	opts[mangos.OptionKeepAlive] = true
	proto := sock.GetProtocol()
	maxrx := 0
	if v, e := sock.GetOption(mangos.OptionMaxRecvSize); e == nil {
		maxrx = v.(int)
	}

	return &dialer{addr: addr, proto: proto, iswss: iswss, opts: opts, maxrx: maxrx}, nil
}

func (t wsTran) NewListener(addr string, sock mangos.Socket) (mangos.PipeListener, error) {
	proto := sock.GetProtocol()
	l, e := t.listener(addr, proto)
	if e == nil {
		if v, e := sock.GetOption(mangos.OptionMaxRecvSize); e == nil {
			l.maxrx = v.(int)
		}
		l.mux.Handle(l.url.Path, l)
	}
	return l, e
}

func (wsTran) listener(addr string, proto mangos.Protocol) (*listener, error) {
	var err error
	l := &listener{proto: proto, opts: make(map[string]interface{})}
	l.cv.L = &l.lock
	l.ug.Subprotocols = []string{proto.Name() + ".sp.nanomsg.org"}

	if strings.HasPrefix(addr, "wss://") {
		l.iswss = true
	}
	l.url, err = url.ParseRequestURI(addr)
	if err != nil {
		return nil, err
	}
	if len(l.url.Path) == 0 {
		l.url.Path = "/"
	}
	l.mux = http.NewServeMux()

	l.htsvr = &http.Server{Addr: l.url.Host, Handler: l.mux}

	return l, nil
}

// NewTransport allocates a new ws:// transport.
func NewTransport() mangos.Transport {
	return wsTran(0)
}
<commit_msg>fixes #262 Set PropTLSConnState for WSS sockets in Dial<commit_after>// Copyright 2016 The Mangos Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use 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 ws implements a simple WebSocket transport for mangos.
// This transport is considered EXPERIMENTAL.
package ws

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/gorilla/websocket"

	"github.com/go-mangos/mangos"
)

// Some special options
const (
	// OptionWebSocketMux is a retrieve-only property used to obtain
	// the *http.ServeMux instance associated with the server.  This
	// can be used to subsequently register additional handlers for
	// different URIs.  This option is only valid on a Listener.
	// Generally you use this option when you want to use the standard
	// mangos Listen() method to start up the server.
	OptionWebSocketMux = "WEBSOCKET-MUX"

	// OptionWebSocketHandler is used to obtain the underlying
	// http.Handler (websocket.Server) object, so you can use this
	// on your own http.Server instances.  It is a gross error to use
	// the value returned by this method on an http server if the
	// server is also started with mangos Listen().  This means that you
	// will use at most either this option, or OptionWebSocketMux, but
	// never both.  This option is only valid on a listener.
	OptionWebSocketHandler = "WEBSOCKET-HANDLER"
)

type options map[string]interface{}

// GetOption retrieves an option value.
func (o options) get(name string) (interface{}, error) {
	if o == nil {
		return nil, mangos.ErrBadOption
	}
	v, ok := o[name]
	if !ok {
		return nil, mangos.ErrBadOption
	}
	return v, nil
}

// SetOption sets an option.  We have none, so just ErrBadOption.
func (o options) set(name string, val interface{}) error {
	switch name {
	case mangos.OptionNoDelay:
		fallthrough
	case mangos.OptionKeepAlive:
		switch v := val.(type) {
		case bool:
			o[name] = v
			return nil
		default:
			return mangos.ErrBadValue
		}
	case mangos.OptionTLSConfig:
		switch v := val.(type) {
		case *tls.Config:
			o[name] = v
			return nil
		default:
			return mangos.ErrBadValue
		}
	}
	return mangos.ErrBadOption
}

// wsPipe implements the Pipe interface on a websocket
type wsPipe struct {
	ws    *websocket.Conn
	proto mangos.Protocol
	addr  string
	open  bool
	wg    sync.WaitGroup
	props map[string]interface{}
	iswss bool
	dtype int
}

type wsTran int

func (w *wsPipe) Recv() (*mangos.Message, error) {

	// We ignore the message type for receive.
	_, body, err := w.ws.ReadMessage()
	if err != nil {
		return nil, err
	}
	msg := mangos.NewMessage(0)
	msg.Body = body
	return msg, nil
}

func (w *wsPipe) Send(m *mangos.Message) error {

	var buf []byte

	if m.Expired() {
		m.Free()
		return nil
	}
	if len(m.Header) > 0 {
		buf = make([]byte, 0, len(m.Header)+len(m.Body))
		buf = append(buf, m.Header...)
		buf = append(buf, m.Body...)
	} else {
		buf = m.Body
	}
	if err := w.ws.WriteMessage(w.dtype, buf); err != nil {
		return err
	}
	m.Free()
	return nil
}

func (w *wsPipe) LocalProtocol() uint16 {
	return w.proto.Number()
}

func (w *wsPipe) RemoteProtocol() uint16 {
	return w.proto.PeerNumber()
}

func (w *wsPipe) Close() error {
	w.open = false
	w.ws.Close()
	w.wg.Done()
	return nil
}

func (w *wsPipe) IsOpen() bool {
	return w.open
}

func (w *wsPipe) GetProp(name string) (interface{}, error) {
	if v, ok := w.props[name]; ok {
		return v, nil
	}
	return nil, mangos.ErrBadProperty
}

type dialer struct {
	addr  string // url
	proto mangos.Protocol
	opts  options
	iswss bool
	maxrx int
}

func (d *dialer) Dial() (mangos.Pipe, error) {
	var w *wsPipe

	wd := &websocket.Dialer{}

	wd.Subprotocols = []string{d.proto.PeerName() + ".sp.nanomsg.org"}
	if v, ok := d.opts[mangos.OptionTLSConfig]; ok {
		wd.TLSClientConfig = v.(*tls.Config)
	}

	w = &wsPipe{proto: d.proto, addr: d.addr, open: true}
	w.dtype = websocket.BinaryMessage
	w.props = make(map[string]interface{})

	var err error
	if w.ws, _, err = wd.Dial(d.addr, nil); err != nil {
		return nil, err
	}
	w.ws.SetReadLimit(int64(d.maxrx))
	w.props[mangos.PropLocalAddr] = w.ws.LocalAddr()
	w.props[mangos.PropRemoteAddr] = w.ws.RemoteAddr()
	if tlsConn, ok := w.ws.UnderlyingConn().(*tls.Conn); ok {
		w.props[mangos.PropTLSConnState] = tlsConn.ConnectionState()
	}

	w.wg.Add(1)
	return w, nil
}

func (d *dialer) SetOption(n string, v interface{}) error {
	return d.opts.set(n, v)
}

func (d *dialer) GetOption(n string) (interface{}, error) {
	return d.opts.get(n)
}

type listener struct {
	pending  []*wsPipe
	lock     sync.Mutex
	cv       sync.Cond
	running  bool
	noserve  bool
	addr     string
	ug       websocket.Upgrader
	htsvr    *http.Server
	mux      *http.ServeMux
	url      *url.URL
	listener net.Listener
	proto    mangos.Protocol
	opts     options
	iswss    bool
	maxrx    int
}

func (l *listener) SetOption(n string, v interface{}) error {
	return l.opts.set(n, v)
}

func (l *listener) GetOption(n string) (interface{}, error) {
	switch n {
	case OptionWebSocketMux:
		return l.mux, nil
	case OptionWebSocketHandler:
		// Caller intends to use use in his own server, so mark
		// us running.  If he didn't mean this, the side effect is
		// that Accept() will appear to hang, even though Listen()
		// is not called yet.
		l.running = true
		l.noserve = true
		return l, nil
	}
	return l.opts.get(n)
}

func (l *listener) Listen() error {
	var taddr *net.TCPAddr
	var err error
	var tcfg *tls.Config

	if l.noserve {
		// The HTTP framework is going to call us, so we use that rather than
		// listening on our own.  We just fake this out.
		return nil
	}
	if l.iswss {
		v, ok := l.opts[mangos.OptionTLSConfig]
		if !ok || v == nil {
			return mangos.ErrTLSNoConfig
		}
		tcfg = v.(*tls.Config)
		if tcfg.Certificates == nil || len(tcfg.Certificates) == 0 {
			return mangos.ErrTLSNoCert
		}
	}

	// We listen separately, that way we can catch and deal with the
	// case of a port already in use.  This also lets us configure
	// properties of the underlying TCP connection.

	if taddr, err = mangos.ResolveTCPAddr(l.url.Host); err != nil {
		return err
	}

	if tlist, err := net.ListenTCP("tcp", taddr); err != nil {
		return err
	} else if l.iswss {
		l.listener = tls.NewListener(tlist, tcfg)
	} else {
		l.listener = tlist
	}
	l.pending = nil
	l.running = true

	l.htsvr = &http.Server{Addr: l.url.Host, Handler: l.mux}

	go l.htsvr.Serve(l.listener)

	return nil
}

func (l *listener) Accept() (mangos.Pipe, error) {
	var w *wsPipe

	l.lock.Lock()
	defer l.lock.Unlock()

	for {
		if !l.running {
			return nil, mangos.ErrClosed
		}
		if len(l.pending) == 0 {
			l.cv.Wait()
			continue
		}
		w = l.pending[len(l.pending)-1]
		l.pending = l.pending[:len(l.pending)-1]
		break
	}

	return w, nil
}

func (l *listener) handler(ws *websocket.Conn, req *http.Request) {
	l.lock.Lock()

	if !l.running {
		ws.Close()
		l.lock.Unlock()
		return
	}

	if ws.Subprotocol() != l.proto.Name()+".sp.nanomsg.org" {
		ws.Close()
		l.lock.Unlock()
		return
	}

	w := &wsPipe{ws: ws, addr: l.addr, proto: l.proto, open: true}
	w.dtype = websocket.BinaryMessage
	w.iswss = l.iswss
	w.ws.SetReadLimit(int64(l.maxrx))

	w.props = make(map[string]interface{})
	w.props[mangos.PropLocalAddr] = ws.LocalAddr()
	w.props[mangos.PropRemoteAddr] = ws.RemoteAddr()

	if req.TLS != nil {
		w.props[mangos.PropTLSConnState] = *req.TLS
	}

	w.wg.Add(1)
	l.pending = append(l.pending, w)
	l.cv.Broadcast()
	l.lock.Unlock()

	// We must not return before the socket is closed, because
	// our caller will close the websocket on our return.
	w.wg.Wait()
}

func (l *listener) Handle(pattern string, handler http.Handler) {
	l.mux.Handle(pattern, handler)
}

func (l *listener) HandleFunc(pattern string, handler http.HandlerFunc) {
	l.mux.HandleFunc(pattern, handler)
}

func (l *listener) Close() error {
	l.lock.Lock()
	defer l.lock.Unlock()
	if !l.running {
		return mangos.ErrClosed
	}
	if l.listener != nil {
		l.listener.Close()
	}
	l.running = false
	l.cv.Broadcast()
	for _, ws := range l.pending {
		ws.Close()
	}
	l.pending = nil
	return nil
}

func (l *listener) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ws, err := l.ug.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	l.handler(ws, r)
}

func (l *listener) Address() string {
	return l.url.String()
}

func (wsTran) Scheme() string {
	return "ws"
}

func (wsTran) NewDialer(addr string, sock mangos.Socket) (mangos.PipeDialer, error) {
	iswss := strings.HasPrefix(addr, "wss://")
	opts := make(map[string]interface{})

	opts[mangos.OptionNoDelay] = true
	opts[mangos.OptionKeepAlive] = true
	proto := sock.GetProtocol()
	maxrx := 0
	if v, e := sock.GetOption(mangos.OptionMaxRecvSize); e == nil {
		maxrx = v.(int)
	}

	return &dialer{addr: addr, proto: proto, iswss: iswss, opts: opts, maxrx: maxrx}, nil
}

func (t wsTran) NewListener(addr string, sock mangos.Socket) (mangos.PipeListener, error) {
	proto := sock.GetProtocol()
	l, e := t.listener(addr, proto)
	if e == nil {
		if v, e := sock.GetOption(mangos.OptionMaxRecvSize); e == nil {
			l.maxrx = v.(int)
		}
		l.mux.Handle(l.url.Path, l)
	}
	return l, e
}

func (wsTran) listener(addr string, proto mangos.Protocol) (*listener, error) {
	var err error
	l := &listener{proto: proto, opts: make(map[string]interface{})}
	l.cv.L = &l.lock
	l.ug.Subprotocols = []string{proto.Name() + ".sp.nanomsg.org"}

	if strings.HasPrefix(addr, "wss://") {
		l.iswss = true
	}
	l.url, err = url.ParseRequestURI(addr)
	if err != nil {
		return nil, err
	}
	if len(l.url.Path) == 0 {
		l.url.Path = "/"
	}
	l.mux = http.NewServeMux()

	l.htsvr = &http.Server{Addr: l.url.Host, Handler: l.mux}

	return l, nil
}

// NewTransport allocates a new ws:// transport.
func NewTransport() mangos.Transport {
	return wsTran(0)
}
<|endoftext|>


<commit_before>package lexer

import (
	"bytes"
	"fmt"
	"github.com/rhysd/gocaml/token"
	"io"
	"unicode"
	"unicode/utf8"
)

type stateFn func(*Lexer) stateFn

const eof = -1

type Lexer struct {
	state   stateFn
	start   token.Position
	current token.Position
	src     *token.Source
	input   *bytes.Reader
	tokens  chan token.Token
	top     rune
	eof     bool
	Error   func(msg string, pos token.Position)
}

func NewLexer(src *token.Source, tokens chan token.Token) *Lexer {
	start := token.Position{
		Offset: 0,
		Line:   1,
		Column: 1,
	}
	return &Lexer{
		state:   lex,
		start:   start,
		current: start,
		input:   bytes.NewReader(src.Code),
		src:     src,
		tokens:  tokens,
		Error:   nil,
	}
}

func (l *Lexer) Lex() {
	// Set top to peek current rune
	l.consume()
	for l.state != nil {
		l.state = l.state(l)
	}
}

func (l *Lexer) emit(kind token.TokenKind) {
	l.tokens <- token.Token{
		kind,
		l.start,
		l.current,
		l.src,
	}
	l.start = l.current
}

func (l *Lexer) emitIdent(ident string) {
	if len(ident) == 1 {
		// Shortcut because no keyword is one character. It must be identifier
		l.emit(token.IDENT)
		return
	}

	switch ident {
	case "true", "false":
		l.emit(token.BOOL)
	case "if":
		l.emit(token.IF)
	case "then":
		l.emit(token.THEN)
	case "else":
		l.emit(token.ELSE)
	case "let":
		l.emit(token.LET)
	case "in":
		l.emit(token.IN)
	case "rec":
		l.emit(token.REC)
	default:
		l.emit(token.IDENT)
	}
}

func (l *Lexer) emitIllegal() {
	t := token.Token{
		token.ILLEGAL,
		l.start,
		l.current,
		l.src,
	}
	l.tokens <- t
	l.start = l.current
}

func (l *Lexer) expected(s string, actual rune) {
	if l.Error != nil {
		l.Error(fmt.Sprintf("Expected %s but got '%c'(%d)", s, actual, actual), l.current)
	}
	l.emitIllegal()
}

func (l *Lexer) unclosedComment(expected string) {
	if l.Error != nil {
		l.Error(fmt.Sprintf("Expected '%s' for closing comment but got EOF", expected), l.current)
	}
	l.emitIllegal()
}

func (l *Lexer) consume() {
	r, _, err := l.input.ReadRune()
	if err == io.EOF {
		l.top = 0
		l.eof = true
		return
	}

	if err != nil {
		panic(err)
	}

	if !utf8.ValidRune(r) {
		panic(fmt.Errorf("Invalid UTF-8 character at line:%d,col:%d: '%c' (%d)", l.current.Line, l.current.Column, r, r))
	}

	l.top = r
	l.eof = false
}

func (l *Lexer) eat() {
	size := utf8.RuneLen(l.top)
	l.current.Offset += size

	// TODO: Consider \n\r
	if l.top == '\n' {
		l.current.Line += 1
		l.current.Column = 1
	} else {
		l.current.Column += size
	}

	l.consume()
}

func (l *Lexer) skip() {
	if l.eof {
		return
	}
	l.eat()
	l.start = l.current
}

func (l *Lexer) eatIndent() bool {
	if !isLetter(l.top) {
		l.expected("letter for head character of identifer", l.top)
		return false
	}
	l.eat()

	for isLetter(l.top) || unicode.IsDigit(l.top) {
		l.eat()
	}
	return true
}

func lexComment(l *Lexer) stateFn {
	for {
		if l.eof {
			l.unclosedComment("*")
			return nil
		}
		if l.top == '*' {
			l.eat()
			if l.eof {
				l.unclosedComment(")")
				return nil
			}
			if l.top == ')' {
				l.eat()
				l.emit(token.COMMENT)
				return lex
			}
		}
		l.eat()
	}
}

func lexLeftParen(l *Lexer) stateFn {
	l.eat()
	if l.top == '*' {
		l.eat()
		return lexComment
	} else {
		l.emit(token.LPAREN)
		return lex
	}
}

func lexAdditiveOp(l *Lexer) stateFn {
	dot, op := token.PLUS_DOT, token.PLUS
	if l.top == '-' {
		dot, op = token.MINUS_DOT, token.MINUS
	}
	l.eat()

	if l.top == '.' {
		l.eat()
		l.emit(dot)
	} else {
		l.emit(op)
	}
	return lex
}

func lexMultOp(l *Lexer) stateFn {
	op := token.STAR_DOT
	if l.top == '/' {
		op = token.SLASH_DOT
	}
	l.eat()

	if l.top != '.' {
		l.expected("'.' for '*.' or '/.' operator", l.top)
		return nil
	}

	l.eat()
	l.emit(op)
	return lex
}

func lexLess(l *Lexer) stateFn {
	l.eat()
	switch l.top {
	case '>':
		l.eat()
		l.emit(token.LESS_GREATER)
	case '=':
		l.eat()
		l.emit(token.LESS_EQUAL)
	case '-':
		l.eat()
		l.emit(token.LESS_MINUS)
	default:
		l.emit(token.LESS)
	}
	return lex
}

func lexGreater(l *Lexer) stateFn {
	l.eat()
	switch l.top {
	case '=':
		l.eat()
		l.emit(token.GREATER_EQUAL)
	default:
		l.emit(token.GREATER)
	}
	return lex
}

// e.g. 123.45e10
func lexNumber(l *Lexer) stateFn {
	tok := token.INT

	// Eat first digit. It's known as digit in lex()
	l.eat()
	for unicode.IsDigit(l.top) {
		l.eat()
	}

	// Note: Allow 1. as 1.0
	if l.top == '.' {
		tok = token.FLOAT
		l.eat()
		for unicode.IsDigit(l.top) {
			l.eat()
		}
	}

	if l.top == 'e' || l.top == 'E' {
		tok = token.FLOAT
		l.eat()
		if l.top == '+' || l.top == '-' {
			l.eat()
		}
		if !unicode.IsDigit(l.top) {
			l.expected("number for exponential part of float literal", l.top)
			return nil
		}
		for unicode.IsDigit(l.top) {
			l.eat()
		}
	}

	l.emit(tok)
	return lex
}

func isLetter(r rune) bool {
	return 'a' <= r && r <= 'z' ||
		'A' <= r && r <= 'Z' ||
		r == '_' ||
		r >= utf8.RuneSelf && unicode.IsLetter(r)
}

func lexArrayCreate(l *Lexer) stateFn {
	if l.top != '.' {
		l.expected("'.' for 'Array.create'", l.top)
		return nil
	}
	l.eat()

	if !l.eatIndent() {
		return nil
	}

	ident := string(l.src.Code[l.start.Offset:l.current.Offset])

	switch ident {
	// Note:
	// Ate 'Array' and '.' but no token was emitted. So 'Array.' remains as
	// current token string.
	case "Array.create", "Array.make":
		l.emit(token.ARRAY_CREATE)
		return lex
	default:
		if l.Error != nil {
			l.Error(fmt.Sprintf("Expected 'create' or 'make' for Array.create but got '%s'", ident), l.current)
		}
		return nil
	}
}

func lexIdent(l *Lexer) stateFn {
	if !l.eatIndent() {
		return nil
	}
	i := string(l.src.Code[l.start.Offset:l.current.Offset])
	if i == "Array" {
		return lexArrayCreate
	}
	l.emitIdent(i)
	return lex
}

func lex(l *Lexer) stateFn {
	for {
		if l.eof {
			l.emit(token.EOF)
			return nil
		}
		switch l.top {
		case '(':
			return lexLeftParen
		case ')':
			l.eat()
			l.emit(token.RPAREN)
		case '+':
			return lexAdditiveOp
		case '-':
			return lexAdditiveOp
		case '*':
			return lexMultOp
		case '/':
			return lexMultOp
		case '=':
			l.eat()
			l.emit(token.EQUAL)
		case '<':
			return lexLess
		case '>':
			return lexGreater
		case ',':
			l.eat()
			l.emit(token.COMMA)
		case '.':
			l.eat()
			l.emit(token.DOT)
		case ';':
			l.eat()
			l.emit(token.SEMICOLON)
		default:
			switch {
			case unicode.IsSpace(l.top):
				l.skip()
			case unicode.IsDigit(l.top):
				return lexNumber
			default:
				return lexIdent
			}
		}
	}
}
<commit_msg>refactorings for lexer<commit_after>package lexer

import (
	"bytes"
	"fmt"
	"github.com/rhysd/gocaml/token"
	"io"
	"unicode"
	"unicode/utf8"
)

type stateFn func(*Lexer) stateFn

const eof = -1

type Lexer struct {
	state   stateFn
	start   token.Position
	current token.Position
	src     *token.Source
	input   *bytes.Reader
	tokens  chan token.Token
	top     rune
	eof     bool
	Error   func(msg string, pos token.Position)
}

func NewLexer(src *token.Source, tokens chan token.Token) *Lexer {
	start := token.Position{
		Offset: 0,
		Line:   1,
		Column: 1,
	}
	return &Lexer{
		state:   lex,
		start:   start,
		current: start,
		input:   bytes.NewReader(src.Code),
		src:     src,
		tokens:  tokens,
		Error:   nil,
	}
}

func (l *Lexer) Lex() {
	// Set top to peek current rune
	l.forward()
	for l.state != nil {
		l.state = l.state(l)
	}
}

func (l *Lexer) emit(kind token.TokenKind) {
	l.tokens <- token.Token{
		kind,
		l.start,
		l.current,
		l.src,
	}
	l.start = l.current
}

func (l *Lexer) emitIdent(ident string) {
	if len(ident) == 1 {
		// Shortcut because no keyword is one character. It must be identifier
		l.emit(token.IDENT)
		return
	}

	switch ident {
	case "true", "false":
		l.emit(token.BOOL)
	case "if":
		l.emit(token.IF)
	case "then":
		l.emit(token.THEN)
	case "else":
		l.emit(token.ELSE)
	case "let":
		l.emit(token.LET)
	case "in":
		l.emit(token.IN)
	case "rec":
		l.emit(token.REC)
	default:
		l.emit(token.IDENT)
	}
}

func (l *Lexer) emitIllegal() {
	t := token.Token{
		token.ILLEGAL,
		l.start,
		l.current,
		l.src,
	}
	l.tokens <- t
	l.start = l.current
}

func (l *Lexer) expected(s string, actual rune) {
	l.errmsg(fmt.Sprintf("Expected %s but got '%c'(%d)", s, actual, actual))
	l.emitIllegal()
}

func (l *Lexer) unclosedComment(expected string) {
	l.errmsg(fmt.Sprintf("Expected '%s' for closing comment but got EOF", expected))
	l.emitIllegal()
}

func (l *Lexer) forward() {
	r, _, err := l.input.ReadRune()
	if err == io.EOF {
		l.top = 0
		l.eof = true
		return
	}

	if err != nil {
		panic(err)
	}

	if !utf8.ValidRune(r) {
		panic(fmt.Errorf("Invalid UTF-8 character at line:%d,col:%d: '%c' (%d)", l.current.Line, l.current.Column, r, r))
	}

	l.top = r
	l.eof = false
}

func (l *Lexer) eat() {
	size := utf8.RuneLen(l.top)
	l.current.Offset += size

	// TODO: Consider \n\r
	if l.top == '\n' {
		l.current.Line += 1
		l.current.Column = 1
	} else {
		l.current.Column += size
	}

	l.forward()
}

func (l *Lexer) consume() {
	if l.eof {
		return
	}
	l.eat()
	l.start = l.current
}

func (l *Lexer) errmsg(msg string) {
	if l.Error == nil {
		return
	}
	l.Error(msg, l.current)
}

func (l *Lexer) eatIndent() bool {
	if !isLetter(l.top) {
		l.expected("letter for head character of identifer", l.top)
		return false
	}
	l.eat()

	for isLetter(l.top) || unicode.IsDigit(l.top) {
		l.eat()
	}
	return true
}

func lexComment(l *Lexer) stateFn {
	for {
		if l.eof {
			l.unclosedComment("*")
			return nil
		}
		if l.top == '*' {
			l.eat()
			if l.eof {
				l.unclosedComment(")")
				return nil
			}
			if l.top == ')' {
				l.eat()
				l.emit(token.COMMENT)
				return lex
			}
		}
		l.eat()
	}
}

func lexLeftParen(l *Lexer) stateFn {
	l.eat()
	if l.top == '*' {
		l.eat()
		return lexComment
	} else {
		l.emit(token.LPAREN)
		return lex
	}
}

func lexAdditiveOp(l *Lexer) stateFn {
	dot, op := token.PLUS_DOT, token.PLUS
	if l.top == '-' {
		dot, op = token.MINUS_DOT, token.MINUS
	}
	l.eat()

	if l.top == '.' {
		l.eat()
		l.emit(dot)
	} else {
		l.emit(op)
	}
	return lex
}

func lexMultOp(l *Lexer) stateFn {
	op := token.STAR_DOT
	if l.top == '/' {
		op = token.SLASH_DOT
	}
	l.eat()

	if l.top != '.' {
		l.expected("'.' for '*.' or '/.' operator", l.top)
		return nil
	}

	l.eat()
	l.emit(op)
	return lex
}

func lexLess(l *Lexer) stateFn {
	l.eat()
	switch l.top {
	case '>':
		l.eat()
		l.emit(token.LESS_GREATER)
	case '=':
		l.eat()
		l.emit(token.LESS_EQUAL)
	case '-':
		l.eat()
		l.emit(token.LESS_MINUS)
	default:
		l.emit(token.LESS)
	}
	return lex
}

func lexGreater(l *Lexer) stateFn {
	l.eat()
	switch l.top {
	case '=':
		l.eat()
		l.emit(token.GREATER_EQUAL)
	default:
		l.emit(token.GREATER)
	}
	return lex
}

// e.g. 123.45e10
func lexNumber(l *Lexer) stateFn {
	tok := token.INT

	// Eat first digit. It's known as digit in lex()
	l.eat()
	for unicode.IsDigit(l.top) {
		l.eat()
	}

	// Note: Allow 1. as 1.0
	if l.top == '.' {
		tok = token.FLOAT
		l.eat()
		for unicode.IsDigit(l.top) {
			l.eat()
		}
	}

	if l.top == 'e' || l.top == 'E' {
		tok = token.FLOAT
		l.eat()
		if l.top == '+' || l.top == '-' {
			l.eat()
		}
		if !unicode.IsDigit(l.top) {
			l.expected("number for exponential part of float literal", l.top)
			return nil
		}
		for unicode.IsDigit(l.top) {
			l.eat()
		}
	}

	l.emit(tok)
	return lex
}

func isLetter(r rune) bool {
	return 'a' <= r && r <= 'z' ||
		'A' <= r && r <= 'Z' ||
		r == '_' ||
		r >= utf8.RuneSelf && unicode.IsLetter(r)
}

func lexArrayCreate(l *Lexer) stateFn {
	if l.top != '.' {
		l.expected("'.' for 'Array.create'", l.top)
		return nil
	}
	l.eat()

	if !l.eatIndent() {
		return nil
	}

	ident := string(l.src.Code[l.start.Offset:l.current.Offset])

	switch ident {
	// Note:
	// Ate 'Array' and '.' but no token was emitted. So 'Array.' remains as
	// current token string.
	case "Array.create", "Array.make":
		l.emit(token.ARRAY_CREATE)
		return lex
	default:
		l.errmsg(fmt.Sprintf("Expected 'create' or 'make' for Array.create but got '%s'", ident))
		return nil
	}
}

func lexIdent(l *Lexer) stateFn {
	if !l.eatIndent() {
		return nil
	}
	i := string(l.src.Code[l.start.Offset:l.current.Offset])
	if i == "Array" {
		return lexArrayCreate
	}
	l.emitIdent(i)
	return lex
}

func lex(l *Lexer) stateFn {
	for {
		if l.eof {
			l.emit(token.EOF)
			return nil
		}
		switch l.top {
		case '(':
			return lexLeftParen
		case ')':
			l.eat()
			l.emit(token.RPAREN)
		case '+':
			return lexAdditiveOp
		case '-':
			return lexAdditiveOp
		case '*':
			return lexMultOp
		case '/':
			return lexMultOp
		case '=':
			l.eat()
			l.emit(token.EQUAL)
		case '<':
			return lexLess
		case '>':
			return lexGreater
		case ',':
			l.eat()
			l.emit(token.COMMA)
		case '.':
			l.eat()
			l.emit(token.DOT)
		case ';':
			l.eat()
			l.emit(token.SEMICOLON)
		default:
			switch {
			case unicode.IsSpace(l.top):
				l.consume()
			case unicode.IsDigit(l.top):
				return lexNumber
			default:
				return lexIdent
			}
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"syscall"

	"github.com/disintegration/imaging"
	"github.com/nfnt/resize"
	"gopkg.in/mgo.v2"
)

// ResizeImageFromGridfs resizes an gridfs image stream
func ResizeImageFromGridfs(originalImage *mgo.GridFile, entry *Entry) (*image.Image, string, error) {
	originalImageData, imageFormat, imgErr := image.Decode(originalImage)

	if imgErr != nil {
		// reset pointer and re read the gridfs image
		originalImage.Seek(0, 0)
		// now it gets really hacky, since go does not support interlaced pngs
		// http://code.google.com/p/go/issues/detail?id=6293
		// we will call imagemagick in order to remove interlacing, and save the image
		// if this fails as well, there must be something wrong
		unresizedImage, magickError := imageMagickFallback(originalImage)

		if magickError == nil {
			return ResizeImage(unresizedImage, imageFormat, entry)
		}

		return nil, imageFormat, imgErr
	}

	return ResizeImage(originalImageData, imageFormat, entry)
}

// imageMagickFallback is used to convert a image with image magick
func imageMagickFallback(originalImage *mgo.GridFile) (image.Image, error) {
	tempDirectory := os.TempDir()

	file, err := ioutil.TempFile(tempDirectory, "magick_original_")
	target, targetErr := ioutil.TempFile(tempDirectory, "magick_target_")

	if targetErr != nil {
		return nil, targetErr
	}

	defer syscall.Unlink(target.Name())

	if err != nil {
		return nil, err
	}

	defer syscall.Unlink(file.Name())

	io.Copy(file, originalImage)
	file.Close()
	originalImage.Close()

	log.Printf("convert %s %s", file.Name(), target.Name())

	cmd := exec.Command("convert", file.Name(), target.Name())

	// blocking execution, since we need to read the image afterwards
	someErr := cmd.Run()

	if someErr != nil {
		return nil, someErr
	}

	targetPNG, openErr := os.Open(target.Name())

	if openErr != nil {
		return nil, openErr
	}

	return png.Decode(targetPNG)
}

// ResizeImage resizes images or crops them if either size is not defined
func ResizeImage(originalImageData image.Image, imageFormat string, entry *Entry) (*image.Image, string, error) {
	if entry.Width < 0 && entry.Height < 0 {
		return nil, "", fmt.Errorf("At least one parameter of width or height must be specified")
	}

	targetHeight := float64(entry.Height)
	targetWidth := float64(entry.Width)

	var dst image.Image
	var err error

	// the Thumbnail method needs correctly adjusted bounds in order to work
	originalBounds := originalImageData.Bounds()
	originalRatio := float64(originalBounds.Dx()) / float64(originalBounds.Dy())

	if entry.Type == TypeResize {
		// the Resize method automatically adjusts ratio based format when one parameter is zero
		if targetWidth < 0 {
			targetWidth = 0
		}

		if targetHeight < 0 {
			targetHeight = 0
		}

		dst = resize.Resize(uint(targetWidth), uint(targetHeight), originalImageData, resize.Lanczos3)
	} else if entry.Type == TypeFit {
		if targetWidth < 0 || targetHeight < 0 {
			return nil, "", fmt.Errorf("When using type fit, both height and width must be specified")
		}

		targetRatio := targetWidth / targetHeight

		if targetRatio < originalRatio {
			targetHeight = targetWidth / originalRatio
		} else {
			targetWidth = targetHeight * originalRatio
		}

		dst = resize.Resize(uint(targetWidth), uint(targetHeight), originalImageData, resize.Lanczos3)
	} else {
		// typeCut
		if targetWidth < 0 {
			targetWidth = float64(targetHeight) * originalRatio
		}

		if targetHeight < 0 {
			targetHeight = float64(targetWidth) / originalRatio
		}

		dst = imaging.Thumbnail(originalImageData, int(targetWidth), int(targetHeight), imaging.Lanczos)
	}

	return &dst, imageFormat, err
}

// EncodeImage encodes the image with the given format
func EncodeImage(targetImage io.Writer, imageData image.Image, imageFormat string) error {
	switch imageFormat {
	case "jpeg":
		jpeg.Encode(targetImage, imageData, &jpeg.Options{jpeg.DefaultQuality})
	case "png":
	    var encoder = png.Encoder
	    encoder.CompressionLevel = png.BestSpeed
		encoder.Encode(targetImage, imageData)
	case "gif":
		gif.Encode(targetImage, imageData, &gif.Options{256, nil, nil})
	default:
		return fmt.Errorf("invalid imageFormat given")
	}

	return nil
}
<commit_msg>Fix wrong usage of encoder<commit_after>package server

import (
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"syscall"

	"github.com/disintegration/imaging"
	"github.com/nfnt/resize"
	"gopkg.in/mgo.v2"
)

// ResizeImageFromGridfs resizes an gridfs image stream
func ResizeImageFromGridfs(originalImage *mgo.GridFile, entry *Entry) (*image.Image, string, error) {
	originalImageData, imageFormat, imgErr := image.Decode(originalImage)

	if imgErr != nil {
		// reset pointer and re read the gridfs image
		originalImage.Seek(0, 0)
		// now it gets really hacky, since go does not support interlaced pngs
		// http://code.google.com/p/go/issues/detail?id=6293
		// we will call imagemagick in order to remove interlacing, and save the image
		// if this fails as well, there must be something wrong
		unresizedImage, magickError := imageMagickFallback(originalImage)

		if magickError == nil {
			return ResizeImage(unresizedImage, imageFormat, entry)
		}

		return nil, imageFormat, imgErr
	}

	return ResizeImage(originalImageData, imageFormat, entry)
}

// imageMagickFallback is used to convert a image with image magick
func imageMagickFallback(originalImage *mgo.GridFile) (image.Image, error) {
	tempDirectory := os.TempDir()

	file, err := ioutil.TempFile(tempDirectory, "magick_original_")
	target, targetErr := ioutil.TempFile(tempDirectory, "magick_target_")

	if targetErr != nil {
		return nil, targetErr
	}

	defer syscall.Unlink(target.Name())

	if err != nil {
		return nil, err
	}

	defer syscall.Unlink(file.Name())

	io.Copy(file, originalImage)
	file.Close()
	originalImage.Close()

	log.Printf("convert %s %s", file.Name(), target.Name())

	cmd := exec.Command("convert", file.Name(), target.Name())

	// blocking execution, since we need to read the image afterwards
	someErr := cmd.Run()

	if someErr != nil {
		return nil, someErr
	}

	targetPNG, openErr := os.Open(target.Name())

	if openErr != nil {
		return nil, openErr
	}

	return png.Decode(targetPNG)
}

// ResizeImage resizes images or crops them if either size is not defined
func ResizeImage(originalImageData image.Image, imageFormat string, entry *Entry) (*image.Image, string, error) {
	if entry.Width < 0 && entry.Height < 0 {
		return nil, "", fmt.Errorf("At least one parameter of width or height must be specified")
	}

	targetHeight := float64(entry.Height)
	targetWidth := float64(entry.Width)

	var dst image.Image
	var err error

	// the Thumbnail method needs correctly adjusted bounds in order to work
	originalBounds := originalImageData.Bounds()
	originalRatio := float64(originalBounds.Dx()) / float64(originalBounds.Dy())

	if entry.Type == TypeResize {
		// the Resize method automatically adjusts ratio based format when one parameter is zero
		if targetWidth < 0 {
			targetWidth = 0
		}

		if targetHeight < 0 {
			targetHeight = 0
		}

		dst = resize.Resize(uint(targetWidth), uint(targetHeight), originalImageData, resize.Lanczos3)
	} else if entry.Type == TypeFit {
		if targetWidth < 0 || targetHeight < 0 {
			return nil, "", fmt.Errorf("When using type fit, both height and width must be specified")
		}

		targetRatio := targetWidth / targetHeight

		if targetRatio < originalRatio {
			targetHeight = targetWidth / originalRatio
		} else {
			targetWidth = targetHeight * originalRatio
		}

		dst = resize.Resize(uint(targetWidth), uint(targetHeight), originalImageData, resize.Lanczos3)
	} else {
		// typeCut
		if targetWidth < 0 {
			targetWidth = float64(targetHeight) * originalRatio
		}

		if targetHeight < 0 {
			targetHeight = float64(targetWidth) / originalRatio
		}

		dst = imaging.Thumbnail(originalImageData, int(targetWidth), int(targetHeight), imaging.Lanczos)
	}

	return &dst, imageFormat, err
}

// EncodeImage encodes the image with the given format
func EncodeImage(targetImage io.Writer, imageData image.Image, imageFormat string) error {
	switch imageFormat {
	case "jpeg":
		jpeg.Encode(targetImage, imageData, &jpeg.Options{jpeg.DefaultQuality})
	case "png":
	    encoder := png.Encoder{CompressionLevel: png.BestSpeed}
		encoder.Encode(targetImage, imageData)
	case "gif":
		gif.Encode(targetImage, imageData, &gif.Options{256, nil, nil})
	default:
		return fmt.Errorf("invalid imageFormat given")
	}

	return nil
}
<|endoftext|>