<commit_before>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
	"sync/atomic"
	"unsafe"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Size of the connection pools (per host).
var PoolSize = 4

// Number of overflow connections allowed in a pool.
var PoolOverflow = PoolSize

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

type VBucketServerMap struct {
	HashAlgorithm string   `json:"hashAlgorithm"`
	NumReplicas   int      `json:"numReplicas"`
	ServerList    []string `json:"serverList"`
	VBucketMap    [][]int  `json:"vBucketMap"`
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	// These are used for JSON IO, but isn't used for processing
	// since it needs to be swapped out safely.
	VBSMJson  VBucketServerMap `json:"vBucketServerMap"`
	NodesJson []Node           `json:"nodes"`

	pool             *Pool
	connPools        unsafe.Pointer // *[]*connectionPool
	vBucketServerMap unsafe.Pointer // *VBucketServerMap
	nodeList         unsafe.Pointer // *[]Node
	commonSufix      string
}

// Get the current vbucket server map
func (b Bucket) VBServerMap() *VBucketServerMap {
	return (*VBucketServerMap)(atomic.LoadPointer(&b.vBucketServerMap))
}

func (b Bucket) Nodes() []Node {
	return *(*[]Node)(atomic.LoadPointer(&b.nodeList))
}

func (b Bucket) getConnPools() []*connectionPool {
	return *(*[]*connectionPool)(atomic.LoadPointer(&b.connPools))
}

func (b *Bucket) replaceConnPools(with []*connectionPool) {
	for {
		old := atomic.LoadPointer(&b.connPools)
		if atomic.CompareAndSwapPointer(&b.connPools, old, unsafe.Pointer(&with)) {
			if old != nil {
				for _, pool := range *(*[]*connectionPool)(old) {
					pool.Close()
				}
			}
			return
		}
	}
}

func (b Bucket) getConnPool(i int) *connectionPool {
	p := b.getConnPools()
	if len(p) > i {
		return p[i]
	}
	return nil
}

func (b Bucket) authHandler() (ah AuthHandler) {
	if b.pool != nil {
		ah = b.pool.client.ah
	}
	if ah == nil {
		ah = &basicAuth{b.Name, ""}
	}
	return
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	vsm := b.VBServerMap()
	rv := make([]string, len(vsm.ServerList))
	copy(rv, vsm.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes() {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	ah       AuthHandler
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, out interface{}) error {
	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, c.ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
		return fmt.Errorf("HTTP error %v getting %q: %s",
			res.Status, u.String(), bod)
	}

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

func basicAuthFromURL(us string) (ah AuthHandler) {
	u, err := ParseURL(us)
	if err != nil {
		return
	}
	if user := u.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}
	return
}

// ConnectWithAuth connects to a couchbase cluster with the given
// authentication handler.
func ConnectWithAuth(baseU string, ah AuthHandler) (c Client, err error) {
	c.BaseURL, err = ParseURL(baseU)
	if err != nil {
		return
	}
	c.ah = ah

	return c, c.parseURLResponse("/pools", &c.Info)
}

// Connect to a couchbase cluster.  An authentication handler will be
// created from the userinfo in the URL if provided.
func Connect(baseU string) (Client, error) {
	return ConnectWithAuth(baseU, basicAuthFromURL(baseU))
}

func (b *Bucket) refresh() error {
	pool := b.pool
	tmpb := &Bucket{}
	err := pool.client.parseURLResponse(b.URI, tmpb)
	if err != nil {
		return err
	}
	newcps := make([]*connectionPool, len(tmpb.VBSMJson.ServerList))
	for i := range newcps {
		newcps[i] = newConnectionPool(
			tmpb.VBSMJson.ServerList[i],
			b.authHandler(), PoolSize, PoolOverflow)
	}
	b.replaceConnPools(newcps)
	atomic.StorePointer(&b.vBucketServerMap, unsafe.Pointer(&tmpb.VBSMJson))
	atomic.StorePointer(&b.nodeList, unsafe.Pointer(&tmpb.NodesJson))
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.nodeList = unsafe.Pointer(&b.NodesJson)
		b.replaceConnPools(make([]*connectionPool, len(b.VBSMJson.ServerList)))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, &p)

	p.client = *c

	err = p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connPools != nil {
		for _, c := range b.getConnPools() {
			if c != nil {
				c.Close()
			}
		}
		b.connPools = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connPools != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	err := rv.refresh()
	if err != nil {
		return nil, err
	}
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<commit_msg>refresh() the new bucket info from a random server<commit_after>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
	"sync/atomic"
	"unsafe"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Size of the connection pools (per host).
var PoolSize = 4

// Number of overflow connections allowed in a pool.
var PoolOverflow = PoolSize

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

type VBucketServerMap struct {
	HashAlgorithm string   `json:"hashAlgorithm"`
	NumReplicas   int      `json:"numReplicas"`
	ServerList    []string `json:"serverList"`
	VBucketMap    [][]int  `json:"vBucketMap"`
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	// These are used for JSON IO, but isn't used for processing
	// since it needs to be swapped out safely.
	VBSMJson  VBucketServerMap `json:"vBucketServerMap"`
	NodesJson []Node           `json:"nodes"`

	pool             *Pool
	connPools        unsafe.Pointer // *[]*connectionPool
	vBucketServerMap unsafe.Pointer // *VBucketServerMap
	nodeList         unsafe.Pointer // *[]Node
	commonSufix      string
}

// Get the current vbucket server map
func (b Bucket) VBServerMap() *VBucketServerMap {
	return (*VBucketServerMap)(atomic.LoadPointer(&b.vBucketServerMap))
}

func (b Bucket) Nodes() []Node {
	return *(*[]Node)(atomic.LoadPointer(&b.nodeList))
}

func (b Bucket) getConnPools() []*connectionPool {
	return *(*[]*connectionPool)(atomic.LoadPointer(&b.connPools))
}

func (b *Bucket) replaceConnPools(with []*connectionPool) {
	for {
		old := atomic.LoadPointer(&b.connPools)
		if atomic.CompareAndSwapPointer(&b.connPools, old, unsafe.Pointer(&with)) {
			if old != nil {
				for _, pool := range *(*[]*connectionPool)(old) {
					pool.Close()
				}
			}
			return
		}
	}
}

func (b Bucket) getConnPool(i int) *connectionPool {
	p := b.getConnPools()
	if len(p) > i {
		return p[i]
	}
	return nil
}

func (b Bucket) authHandler() (ah AuthHandler) {
	if b.pool != nil {
		ah = b.pool.client.ah
	}
	if ah == nil {
		ah = &basicAuth{b.Name, ""}
	}
	return
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	vsm := b.VBServerMap()
	rv := make([]string, len(vsm.ServerList))
	copy(rv, vsm.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes() {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	ah       AuthHandler
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func queryRestAPI(
	baseUrl *url.URL,
	path string,
	authHandler AuthHandler,
	out interface{}) error {
	u := *baseUrl
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, authHandler)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
		return fmt.Errorf("HTTP error %v getting %q: %s",
			res.Status, u.String(), bod)
	}

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

func (c *Client) parseURLResponse(path string, out interface{}) error {
	return queryRestAPI(c.BaseURL, path, c.ah, out)
}

func (b *Bucket) parseURLResponse(path string, out interface{}) error {
	nodes := b.Nodes()
	if len(nodes) == 0 {
		return errors.New("no couch rest URLs")
	}

	// Pick a random node to start querying.
	startNode := rand.Intn(len(nodes))
	maxRetries := len(nodes)
	for i := 0; i < maxRetries; i++ {
		node := nodes[(startNode+i)%len(nodes)] // Wrap around the nodes list.
		// Skip non-healthy nodes.
		if node.Status != "healthy" {
			continue
		}

		url := &url.URL{
			Host:   node.Hostname,
			Scheme: "http",
		}

		err := queryRestAPI(url, path, b.pool.client.ah, out)
		if err == nil {
			return err
		}
	}
	return errors.New("All nodes failed to respond.")
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

func basicAuthFromURL(us string) (ah AuthHandler) {
	u, err := ParseURL(us)
	if err != nil {
		return
	}
	if user := u.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}
	return
}

// ConnectWithAuth connects to a couchbase cluster with the given
// authentication handler.
func ConnectWithAuth(baseU string, ah AuthHandler) (c Client, err error) {
	c.BaseURL, err = ParseURL(baseU)
	if err != nil {
		return
	}
	c.ah = ah

	return c, c.parseURLResponse("/pools", &c.Info)
}

// Connect to a couchbase cluster.  An authentication handler will be
// created from the userinfo in the URL if provided.
func Connect(baseU string) (Client, error) {
	return ConnectWithAuth(baseU, basicAuthFromURL(baseU))
}

func (b *Bucket) refresh() error {
	pool := b.pool
	tmpb := &Bucket{}
	err := pool.client.parseURLResponse(b.URI, tmpb)
	if err != nil {
		return err
	}
	newcps := make([]*connectionPool, len(tmpb.VBSMJson.ServerList))
	for i := range newcps {
		newcps[i] = newConnectionPool(
			tmpb.VBSMJson.ServerList[i],
			b.authHandler(), PoolSize, PoolOverflow)
	}
	b.replaceConnPools(newcps)
	atomic.StorePointer(&b.vBucketServerMap, unsafe.Pointer(&tmpb.VBSMJson))
	atomic.StorePointer(&b.nodeList, unsafe.Pointer(&tmpb.NodesJson))
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.nodeList = unsafe.Pointer(&b.NodesJson)
		b.replaceConnPools(make([]*connectionPool, len(b.VBSMJson.ServerList)))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, &p)

	p.client = *c

	err = p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connPools != nil {
		for _, c := range b.getConnPools() {
			if c != nil {
				c.Close()
			}
		}
		b.connPools = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connPools != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	err := rv.refresh()
	if err != nil {
		return nil, err
	}
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

type SocketInfo struct {
	Src_ip, Dst_ip     uint32
	Src_port, Dst_port uint16

	Uid   uint16
	Inode int64
}

type PortProcMapping struct {
	Port uint16
	Pid  int
	Proc *Process
}

type Process struct {
	Name    string
	Grepper string
	Pids    []int

	proc *ProcessesWatcher

	RefreshPidsTimer <-chan time.Time
}

type ProcessesWatcher struct {
	PortProcMap   map[uint16]PortProcMapping
	LastMapUpdate time.Time
	Processes     []*Process
	LocalAddrs    []net.IP

	// config
	ReadFromProc    bool
	MaxReadFreq     time.Duration
	RefreshPidsFreq time.Duration

	// test helpers
	proc_prefix string
	TestSignals *chan bool
}

var procWatcher ProcessesWatcher

// Config
type tomlProcs struct {
	Dont_read_from_proc bool
	Max_proc_read_freq  int
	Monitored           map[string]tomlProc
	Refresh_pids_freq   int
}

type tomlProc struct {
	Cmdline_grep string
}

func (proc *ProcessesWatcher) Init(config *tomlProcs) error {

	proc.proc_prefix = ""
	proc.PortProcMap = make(map[uint16]PortProcMapping)
	proc.LastMapUpdate = time.Now()

	proc.ReadFromProc = !config.Dont_read_from_proc
	if proc.ReadFromProc {
		if runtime.GOOS != "linux" {
			proc.ReadFromProc = false
			INFO("Disabled /proc/ reading because not on linux")
		}
	}

	if config.Max_proc_read_freq == 0 {
		proc.MaxReadFreq = 10 * time.Millisecond
	} else {
		proc.MaxReadFreq = time.Duration(config.Max_proc_read_freq) * time.Millisecond
	}

	if config.Refresh_pids_freq == 0 {
		proc.RefreshPidsFreq = 1 * time.Second
	} else {
		proc.RefreshPidsFreq = time.Duration(config.Refresh_pids_freq) * time.Millisecond
	}

	// Read the local IP addresses
	var err error
	proc.LocalAddrs, err = LocalIpAddrs()
	if err != nil {
		ERR("Error getting local IP addresses: %s", err)
		proc.LocalAddrs = []net.IP{}
	}

	if proc.ReadFromProc {
		for pstr, procConfig := range config.Monitored {

			grepper := procConfig.Cmdline_grep
			if len(grepper) == 0 {
				grepper = pstr
			}

			p, err := NewProcess(proc, pstr, grepper, time.Tick(proc.RefreshPidsFreq))
			if err != nil {
				ERR("NewProcess: %s", err)
			} else {
				proc.Processes = append(proc.Processes, p)
			}
		}
	}

	return nil
}

func NewProcess(proc *ProcessesWatcher, name string, grepper string,
	refreshPidsTimer <-chan time.Time) (*Process, error) {

	p := &Process{Name: name, proc: proc, Grepper: grepper,
		RefreshPidsTimer: refreshPidsTimer}

	// start periodic timer in its own goroutine
	go p.RefreshPids()

	return p, nil
}

func (p *Process) RefreshPids() {
	DEBUG("procs", "In RefreshPids")
	for _ = range p.RefreshPidsTimer {
		DEBUG("procs", "In RefreshPids tick")
		var err error
		p.Pids, err = FindPidsByCmdlineGrep(p.proc.proc_prefix, p.Grepper)
		if err != nil {
			ERR("Error finding PID files for %s: %s", p.Name, err)
		}
		DEBUG("procs", "RefreshPids found pids %s for process %s", p.Pids, p.Name)

		if p.proc.TestSignals != nil {
			*p.proc.TestSignals <- true
		}
	}
}

func FindPidsByCmdlineGrep(prefix string, process string) ([]int, error) {
	defer RECOVER("FindPidsByCmdlineGrep exception")
	pids := []int{}

	proc, err := os.Open(filepath.Join(prefix, "/proc"))
	if err != nil {
		return pids, MsgError("Open /proc: %s", err)
	}

	names, err := proc.Readdirnames(0)
	if err != nil {
		return pids, MsgError("Readdirnames: %s", err)
	}

	for _, name := range names {
		pid, err := strconv.Atoi(name)
		if err != nil {
			continue
		}

		cmdline, err := ioutil.ReadFile(filepath.Join(prefix, "/proc/", name, "cmdline"))
		if err != nil {
			continue
		}

		if strings.Index(string(cmdline), process) >= 0 {
			pids = append(pids, pid)
		}
	}

	return pids, nil
}

func (proc *ProcessesWatcher) FindProcessesTuple(tuple *IpPortTuple) (proc_tuple *CmdlineTuple) {
	proc_tuple = &CmdlineTuple{}

	if !proc.ReadFromProc {
		return
	}

	if proc.IsLocalIp(tuple.Src_ip) {
		DEBUG("procs", "Looking for port %d", tuple.Src_port)
		proc_tuple.Src = []byte(proc.FindProc(tuple.Src_port))
		if len(proc_tuple.Src) > 0 {
			DEBUG("procs", "Found device %s for port %d", proc_tuple.Src, tuple.Src_port)
		}
	}

	if proc.IsLocalIp(tuple.Dst_ip) {
		DEBUG("procs", "Looking for port %d", tuple.Dst_port)
		proc_tuple.Dst = []byte(proc.FindProc(tuple.Dst_port))
		if len(proc_tuple.Dst) > 0 {
			DEBUG("procs", "Found device %s for port %d", proc_tuple.Dst, tuple.Dst_port)
		}
	}

	return
}

func (proc *ProcessesWatcher) FindProc(port uint16) (procname string) {
	procname = ""
	defer RECOVER("FindProc exception")

	p, exists := proc.PortProcMap[port]
	if exists {
		return p.Proc.Name
	}

	if time.Now().Sub(proc.LastMapUpdate) > proc.MaxReadFreq {
		proc.UpdateMap()

		// try again
		p, exists := proc.PortProcMap[port]
		if exists {
			return p.Proc.Name
		}
	}

	return ""
}

func hex_to_ip_port(str []byte) (uint32, uint16, error) {
	words := bytes.Split(str, []byte(":"))
	if len(words) < 2 {
		return 0, 0, MsgError("Didn't find ':' as a separator")
	}

	ip, err := strconv.ParseInt(string(words[0]), 16, 64)
	if err != nil {
		return 0, 0, err
	}

	port, err := strconv.ParseInt(string(words[1]), 16, 32)
	if err != nil {
		return 0, 0, err
	}

	return uint32(ip), uint16(port), nil
}

func (proc *ProcessesWatcher) UpdateMap() {

	DEBUG("procs", "UpdateMap()")
	file, err := os.Open("/proc/net/tcp")
	if err != nil {
		ERR("Open: %s", err)
		return
	}
	socks, err := Parse_Proc_Net_Tcp(file)
	if err != nil {
		ERR("Parse_Proc_Net_Tcp: %s", err)
		return
	}
	socks_map := map[int64]*SocketInfo{}
	for _, s := range socks {
		socks_map[s.Inode] = s
	}

	for _, p := range proc.Processes {
		for _, pid := range p.Pids {
			inodes, err := FindSocketsOfPid(proc.proc_prefix, pid)
			if err != nil {
				ERR("FindSocketsOfPid: %s", err)
				continue
			}

			for _, inode := range inodes {
				sockInfo, exists := socks_map[inode]
				if exists {
					proc.UpdateMappingEntry(sockInfo.Src_port, pid, p)
				}
			}

		}
	}
}

// Parses the /proc/net/tcp file
func Parse_Proc_Net_Tcp(input io.Reader) ([]*SocketInfo, error) {
	buf := bufio.NewReader(input)

	sockets := []*SocketInfo{}
	var err error = nil
	var line []byte
	for err != io.EOF {
		line, err = buf.ReadBytes('\n')
		if err != nil && err != io.EOF {
			ERR("Error reading /proc/net/tcp: %s", err)
			return nil, err
		}

		words := bytes.Fields(line)
		if len(words) < 10 || bytes.Equal(words[0], []byte("sl")) {
			DEBUG("procs", "Less then 10 words (%d) or starting with 'sl': %s", len(words), words)
			continue
		}

		var sock SocketInfo
		var err_ error

		sock.Src_ip, sock.Src_port, err_ = hex_to_ip_port(words[1])
		if err_ != nil {
			DEBUG("procs", "Error parsing IP and port: %s", err_)
			continue
		}

		sock.Dst_ip, sock.Dst_port, err_ = hex_to_ip_port(words[2])
		if err_ != nil {
			DEBUG("procs", "Error parsing IP and port: %s", err_)
			continue
		}

		uid, _ := strconv.Atoi(string(words[7]))
		sock.Uid = uint16(uid)
		inode, _ := strconv.Atoi(string(words[9]))
		sock.Inode = int64(inode)

		sockets = append(sockets, &sock)
	}
	return sockets, nil
}

func (proc *ProcessesWatcher) UpdateMappingEntry(port uint16, pid int, p *Process) {
	entry := PortProcMapping{Port: port, Pid: pid, Proc: p}

	// Simply overwrite old entries for now.
	// We never expire entries from this map. Since there are 65k possible
	// ports, the size of the dict can be max 1.5 MB, which we consider
	// reasonable.
	proc.PortProcMap[port] = entry

	DEBUG("procsdetailed", "UpdateMappingEntry(): port=%d pid=%d", port, p.Name)
}

func FindSocketsOfPid(prefix string, pid int) (inodes []int64, err error) {

	dirname := filepath.Join(prefix, "/proc", strconv.Itoa(pid), "fd")
	procfs, err := os.Open(dirname)
	if err != nil {
		return []int64{}, MsgError("Open: %s", err)
	}
	names, err := procfs.Readdirnames(0)
	if err != nil {
		return []int64{}, MsgError("Readdirnames: %s", err)
	}

	for _, name := range names {
		link, err := os.Readlink(filepath.Join(dirname, name))
		if err != nil {
			DEBUG("procs", "Readlink %s: %s", name, err)
			continue
		}

		if strings.HasPrefix(link, "socket:[") {
			inode, err := strconv.ParseInt(link[8:len(link)-1], 10, 64)
			if err != nil {
				DEBUG("procs", "ParseInt: %s:", err)
				continue
			}

			inodes = append(inodes, int64(inode))
		}
	}

	return inodes, nil
}

func (proc *ProcessesWatcher) IsLocalIp(ip net.IP) bool {

	if ip.IsLoopback() {
		return true
	}

	for _, addr := range proc.LocalAddrs {
		if ip.Equal(addr) {
			return true
		}
	}

	return false
}
<commit_msg>Bug fix: update the LastMapUpdate<commit_after>package main

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

type SocketInfo struct {
	Src_ip, Dst_ip     uint32
	Src_port, Dst_port uint16

	Uid   uint16
	Inode int64
}

type PortProcMapping struct {
	Port uint16
	Pid  int
	Proc *Process
}

type Process struct {
	Name    string
	Grepper string
	Pids    []int

	proc *ProcessesWatcher

	RefreshPidsTimer <-chan time.Time
}

type ProcessesWatcher struct {
	PortProcMap   map[uint16]PortProcMapping
	LastMapUpdate time.Time
	Processes     []*Process
	LocalAddrs    []net.IP

	// config
	ReadFromProc    bool
	MaxReadFreq     time.Duration
	RefreshPidsFreq time.Duration

	// test helpers
	proc_prefix string
	TestSignals *chan bool
}

var procWatcher ProcessesWatcher

// Config
type tomlProcs struct {
	Dont_read_from_proc bool
	Max_proc_read_freq  int
	Monitored           map[string]tomlProc
	Refresh_pids_freq   int
}

type tomlProc struct {
	Cmdline_grep string
}

func (proc *ProcessesWatcher) Init(config *tomlProcs) error {

	proc.proc_prefix = ""
	proc.PortProcMap = make(map[uint16]PortProcMapping)
	proc.LastMapUpdate = time.Now()

	proc.ReadFromProc = !config.Dont_read_from_proc
	if proc.ReadFromProc {
		if runtime.GOOS != "linux" {
			proc.ReadFromProc = false
			INFO("Disabled /proc/ reading because not on linux")
		}
	}

	if config.Max_proc_read_freq == 0 {
		proc.MaxReadFreq = 10 * time.Millisecond
	} else {
		proc.MaxReadFreq = time.Duration(config.Max_proc_read_freq) * time.Millisecond
	}

	if config.Refresh_pids_freq == 0 {
		proc.RefreshPidsFreq = 1 * time.Second
	} else {
		proc.RefreshPidsFreq = time.Duration(config.Refresh_pids_freq) * time.Millisecond
	}

	// Read the local IP addresses
	var err error
	proc.LocalAddrs, err = LocalIpAddrs()
	if err != nil {
		ERR("Error getting local IP addresses: %s", err)
		proc.LocalAddrs = []net.IP{}
	}

	if proc.ReadFromProc {
		for pstr, procConfig := range config.Monitored {

			grepper := procConfig.Cmdline_grep
			if len(grepper) == 0 {
				grepper = pstr
			}

			p, err := NewProcess(proc, pstr, grepper, time.Tick(proc.RefreshPidsFreq))
			if err != nil {
				ERR("NewProcess: %s", err)
			} else {
				proc.Processes = append(proc.Processes, p)
			}
		}
	}

	return nil
}

func NewProcess(proc *ProcessesWatcher, name string, grepper string,
	refreshPidsTimer <-chan time.Time) (*Process, error) {

	p := &Process{Name: name, proc: proc, Grepper: grepper,
		RefreshPidsTimer: refreshPidsTimer}

	// start periodic timer in its own goroutine
	go p.RefreshPids()

	return p, nil
}

func (p *Process) RefreshPids() {
	DEBUG("procs", "In RefreshPids")
	for _ = range p.RefreshPidsTimer {
		DEBUG("procs", "In RefreshPids tick")
		var err error
		p.Pids, err = FindPidsByCmdlineGrep(p.proc.proc_prefix, p.Grepper)
		if err != nil {
			ERR("Error finding PID files for %s: %s", p.Name, err)
		}
		DEBUG("procs", "RefreshPids found pids %s for process %s", p.Pids, p.Name)

		if p.proc.TestSignals != nil {
			*p.proc.TestSignals <- true
		}
	}
}

func FindPidsByCmdlineGrep(prefix string, process string) ([]int, error) {
	defer RECOVER("FindPidsByCmdlineGrep exception")
	pids := []int{}

	proc, err := os.Open(filepath.Join(prefix, "/proc"))
	if err != nil {
		return pids, MsgError("Open /proc: %s", err)
	}

	names, err := proc.Readdirnames(0)
	if err != nil {
		return pids, MsgError("Readdirnames: %s", err)
	}

	for _, name := range names {
		pid, err := strconv.Atoi(name)
		if err != nil {
			continue
		}

		cmdline, err := ioutil.ReadFile(filepath.Join(prefix, "/proc/", name, "cmdline"))
		if err != nil {
			continue
		}

		if strings.Index(string(cmdline), process) >= 0 {
			pids = append(pids, pid)
		}
	}

	return pids, nil
}

func (proc *ProcessesWatcher) FindProcessesTuple(tuple *IpPortTuple) (proc_tuple *CmdlineTuple) {
	proc_tuple = &CmdlineTuple{}

	if !proc.ReadFromProc {
		return
	}

	if proc.IsLocalIp(tuple.Src_ip) {
		DEBUG("procs", "Looking for port %d", tuple.Src_port)
		proc_tuple.Src = []byte(proc.FindProc(tuple.Src_port))
		if len(proc_tuple.Src) > 0 {
			DEBUG("procs", "Found device %s for port %d", proc_tuple.Src, tuple.Src_port)
		}
	}

	if proc.IsLocalIp(tuple.Dst_ip) {
		DEBUG("procs", "Looking for port %d", tuple.Dst_port)
		proc_tuple.Dst = []byte(proc.FindProc(tuple.Dst_port))
		if len(proc_tuple.Dst) > 0 {
			DEBUG("procs", "Found device %s for port %d", proc_tuple.Dst, tuple.Dst_port)
		}
	}

	return
}

func (proc *ProcessesWatcher) FindProc(port uint16) (procname string) {
	procname = ""
	defer RECOVER("FindProc exception")

	p, exists := proc.PortProcMap[port]
	if exists {
		return p.Proc.Name
	}

	if time.Now().Sub(proc.LastMapUpdate) > proc.MaxReadFreq {
		proc.UpdateMap()

		// try again
		p, exists := proc.PortProcMap[port]
		if exists {
			return p.Proc.Name
		}
	}

	return ""
}

func hex_to_ip_port(str []byte) (uint32, uint16, error) {
	words := bytes.Split(str, []byte(":"))
	if len(words) < 2 {
		return 0, 0, MsgError("Didn't find ':' as a separator")
	}

	ip, err := strconv.ParseInt(string(words[0]), 16, 64)
	if err != nil {
		return 0, 0, err
	}

	port, err := strconv.ParseInt(string(words[1]), 16, 32)
	if err != nil {
		return 0, 0, err
	}

	return uint32(ip), uint16(port), nil
}

func (proc *ProcessesWatcher) UpdateMap() {

	DEBUG("procs", "UpdateMap()")
	file, err := os.Open("/proc/net/tcp")
	if err != nil {
		ERR("Open: %s", err)
		return
	}
	socks, err := Parse_Proc_Net_Tcp(file)
	if err != nil {
		ERR("Parse_Proc_Net_Tcp: %s", err)
		return
	}
	socks_map := map[int64]*SocketInfo{}
	for _, s := range socks {
		socks_map[s.Inode] = s
	}

	for _, p := range proc.Processes {
		for _, pid := range p.Pids {
			inodes, err := FindSocketsOfPid(proc.proc_prefix, pid)
			if err != nil {
				ERR("FindSocketsOfPid: %s", err)
				continue
			}

			for _, inode := range inodes {
				sockInfo, exists := socks_map[inode]
				if exists {
					proc.UpdateMappingEntry(sockInfo.Src_port, pid, p)
				}
			}

		}
	}

	proc.LastMapUpdate = time.Now()
}

// Parses the /proc/net/tcp file
func Parse_Proc_Net_Tcp(input io.Reader) ([]*SocketInfo, error) {
	buf := bufio.NewReader(input)

	sockets := []*SocketInfo{}
	var err error = nil
	var line []byte
	for err != io.EOF {
		line, err = buf.ReadBytes('\n')
		if err != nil && err != io.EOF {
			ERR("Error reading /proc/net/tcp: %s", err)
			return nil, err
		}

		words := bytes.Fields(line)
		if len(words) < 10 || bytes.Equal(words[0], []byte("sl")) {
			DEBUG("procs", "Less then 10 words (%d) or starting with 'sl': %s", len(words), words)
			continue
		}

		var sock SocketInfo
		var err_ error

		sock.Src_ip, sock.Src_port, err_ = hex_to_ip_port(words[1])
		if err_ != nil {
			DEBUG("procs", "Error parsing IP and port: %s", err_)
			continue
		}

		sock.Dst_ip, sock.Dst_port, err_ = hex_to_ip_port(words[2])
		if err_ != nil {
			DEBUG("procs", "Error parsing IP and port: %s", err_)
			continue
		}

		uid, _ := strconv.Atoi(string(words[7]))
		sock.Uid = uint16(uid)
		inode, _ := strconv.Atoi(string(words[9]))
		sock.Inode = int64(inode)

		sockets = append(sockets, &sock)
	}
	return sockets, nil
}

func (proc *ProcessesWatcher) UpdateMappingEntry(port uint16, pid int, p *Process) {
	entry := PortProcMapping{Port: port, Pid: pid, Proc: p}

	// Simply overwrite old entries for now.
	// We never expire entries from this map. Since there are 65k possible
	// ports, the size of the dict can be max 1.5 MB, which we consider
	// reasonable.
	proc.PortProcMap[port] = entry

	DEBUG("procsdetailed", "UpdateMappingEntry(): port=%d pid=%d", port, p.Name)
}

func FindSocketsOfPid(prefix string, pid int) (inodes []int64, err error) {

	dirname := filepath.Join(prefix, "/proc", strconv.Itoa(pid), "fd")
	procfs, err := os.Open(dirname)
	if err != nil {
		return []int64{}, MsgError("Open: %s", err)
	}
	names, err := procfs.Readdirnames(0)
	if err != nil {
		return []int64{}, MsgError("Readdirnames: %s", err)
	}

	for _, name := range names {
		link, err := os.Readlink(filepath.Join(dirname, name))
		if err != nil {
			DEBUG("procs", "Readlink %s: %s", name, err)
			continue
		}

		if strings.HasPrefix(link, "socket:[") {
			inode, err := strconv.ParseInt(link[8:len(link)-1], 10, 64)
			if err != nil {
				DEBUG("procs", "ParseInt: %s:", err)
				continue
			}

			inodes = append(inodes, int64(inode))
		}
	}

	return inodes, nil
}

func (proc *ProcessesWatcher) IsLocalIp(ip net.IP) bool {

	if ip.IsLoopback() {
		return true
	}

	for _, addr := range proc.LocalAddrs {
		if ip.Equal(addr) {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2015 YAMAMOTO Masaya <pandax381@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */
package main

import (
	"crypto/tls"
	"flag"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

type config struct {
	laddr      string
	saddr      string
	raddr      string
	tlsAccept  bool
	tlsKey     string
	tlsCert    string
	tlsConnect bool
}

type proxy struct {
	conf config
	quit chan bool
}

func NewProxy(conf config) *proxy {
	return &proxy{
		conf: conf,
		quit: make(chan bool),
	}
}

func (p *proxy) Run() {
	listener, err := net.Listen("tcp", p.conf.laddr)
	if err != nil {
		log.Println(err)
		return
	}
	log.Printf("Listen on %s\n", listener.Addr())
	complete := make(chan bool)
    go func() {
        wg := &sync.WaitGroup{}
        quit := make(chan bool)
        for {
            conn, err := listener.Accept()
            if err != nil {
                if e, ok := err.(net.Error); ok && e.Temporary() {
                    continue
                }
                listener.Close()
                close(quit)
                wg.Wait()
                close(complete)
                return
            }
            wg.Add(1)
            go p.handle(conn, wg, quit)
        }
    }()
	for {
		select {
		case <-p.quit:
			listener.Close()
		case <-complete:
			return
		}
	}
	// does not reach
}

func (p *proxy) handle(conn1 net.Conn, wg *sync.WaitGroup, quit chan bool) {
    defer wg.Done()
    defer log.Println("Close Session")
    log.Println("Accept New Session")
    if p.conf.tlsAccept {
        cert, err := tls.LoadX509KeyPair(p.conf.tlsCert, p.conf.tlsKey)
        if err != nil {
            log.Println(err)
            conn1.Close()
            return
        }
        tlsConfig := &tls.Config{
            Certificates: []tls.Certificate{cert},
        }
        conn1 = tls.Server(conn1, tlsConfig)
    }
    log.Println("Connect Remote Host")
    saddr, err := net.ResolveTCPAddr("tcp", p.conf.saddr)
    if err != nil {
        log.Println(err)
        conn1.Close()
        return
    }
    raddr, err := net.ResolveTCPAddr("tcp", p.conf.raddr)
    if err != nil {
        log.Println(err)
        conn1.Close()
        return
    }
    var conn2 net.Conn
    conn2, err = net.DialTCP("tcp", saddr, raddr)
    if err != nil {
        log.Println(err)
        conn1.Close()
        return
    }
    if p.conf.tlsConnect {
        tlsConfig := &tls.Config{
            InsecureSkipVerify: true,
        }
        conn2 = tls.Client(conn2, tlsConfig)
    }
    complete := make(chan int64)
    go transfer(conn1, conn2, complete)
    go transfer(conn2, conn1, complete)
    for n := 2; n > 0; n-- {
        select {
        case <-complete:
            break
        case <-quit:
            conn1.Close()
            conn2.Close()
            for ; n > 0; n-- {
                <-complete
            }
            return
        }
    }
    // just in case
    conn1.Close()
    conn2.Close()
}

func transfer(dst, src net.Conn, complete chan<- int64) {
    n, err := io.Copy(dst, src)
    if err != nil {
        log.Println(err)
        if e, ok := err.(*net.OpError); ok && e.Err == syscall.EPIPE {
            if _, ok := src.(*net.TCPConn); ok {
                src.(*net.TCPConn).CloseRead()
            } else {
                src.Close()
            }
        }
    }
    if _, ok := dst.(*net.TCPConn); ok {
        dst.(*net.TCPConn).CloseWrite()
    } else {
        dst.Close()
    }
    complete <- n
}

func (p *proxy) shutdown() {
	close(p.quit)
}

func main() {
	laddr := flag.String("l", ":8000", "Listen Address")
	saddr := flag.String("s", "", "Source Address")
	raddr := flag.String("r", "localhost:8080", "Remote Address")
	tlsAccept := flag.Bool("tls-accept", false, "Enable TLS Accept")
	tlsCert := flag.String("tls-cert", "./server.crt", "Certificate File")
	tlsKey := flag.String("tls-key", "./server.key", "Privatekey File")
	tlsConnect := flag.Bool("tls-connect", false, "Enable TLS Connect")
	flag.Parse()

	log.SetFlags(log.Ldate | log.Lmicroseconds | log.Lshortfile)
	log.Println("PID:", os.Getpid())

	config := config{
		laddr:      *laddr,
		saddr:      *saddr,
		raddr:      *raddr,
		tlsAccept:  *tlsAccept,
		tlsCert:    *tlsCert,
		tlsKey:     *tlsKey,
		tlsConnect: *tlsConnect,
	}
	proxy := NewProxy(config)
	complete := make(chan bool)
	go func() {
		defer close(complete)
		proxy.Run()
	}()
	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt)
	for {
		select {
		case s := <-sigch:
			log.Println("Recieve signal:", s)
			proxy.shutdown()
		case <-complete:
			log.Println("Good bye")
			return
		}
	}
	// does not reach
}
<commit_msg>go fmt かけてなかったのでかけた<commit_after>/*
 * Copyright (c) 2015 YAMAMOTO Masaya <pandax381@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */
package main

import (
	"crypto/tls"
	"flag"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

type config struct {
	laddr      string
	saddr      string
	raddr      string
	tlsAccept  bool
	tlsKey     string
	tlsCert    string
	tlsConnect bool
}

type proxy struct {
	conf config
	quit chan bool
}

func NewProxy(conf config) *proxy {
	return &proxy{
		conf: conf,
		quit: make(chan bool),
	}
}

func (p *proxy) Run() {
	listener, err := net.Listen("tcp", p.conf.laddr)
	if err != nil {
		log.Println(err)
		return
	}
	log.Printf("Listen on %s\n", listener.Addr())
	complete := make(chan bool)
	go func() {
		wg := &sync.WaitGroup{}
		quit := make(chan bool)
		for {
			conn, err := listener.Accept()
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Temporary() {
					continue
				}
				listener.Close()
				close(quit)
				wg.Wait()
				close(complete)
				return
			}
			wg.Add(1)
			go p.handle(conn, wg, quit)
		}
	}()
	for {
		select {
		case <-p.quit:
			listener.Close()
		case <-complete:
			return
		}
	}
	// does not reach
}

func (p *proxy) handle(conn1 net.Conn, wg *sync.WaitGroup, quit chan bool) {
	defer wg.Done()
	defer log.Println("Close Session")
	log.Println("Accept New Session")
	if p.conf.tlsAccept {
		cert, err := tls.LoadX509KeyPair(p.conf.tlsCert, p.conf.tlsKey)
		if err != nil {
			log.Println(err)
			conn1.Close()
			return
		}
		tlsConfig := &tls.Config{
			Certificates: []tls.Certificate{cert},
		}
		conn1 = tls.Server(conn1, tlsConfig)
	}
	log.Println("Connect Remote Host")
	saddr, err := net.ResolveTCPAddr("tcp", p.conf.saddr)
	if err != nil {
		log.Println(err)
		conn1.Close()
		return
	}
	raddr, err := net.ResolveTCPAddr("tcp", p.conf.raddr)
	if err != nil {
		log.Println(err)
		conn1.Close()
		return
	}
	var conn2 net.Conn
	conn2, err = net.DialTCP("tcp", saddr, raddr)
	if err != nil {
		log.Println(err)
		conn1.Close()
		return
	}
	if p.conf.tlsConnect {
		tlsConfig := &tls.Config{
			InsecureSkipVerify: true,
		}
		conn2 = tls.Client(conn2, tlsConfig)
	}
	complete := make(chan int64)
	go transfer(conn1, conn2, complete)
	go transfer(conn2, conn1, complete)
	for n := 2; n > 0; n-- {
		select {
		case <-complete:
			break
		case <-quit:
			conn1.Close()
			conn2.Close()
			for ; n > 0; n-- {
				<-complete
			}
			return
		}
	}
	// just in case
	conn1.Close()
	conn2.Close()
}

func transfer(dst, src net.Conn, complete chan<- int64) {
	n, err := io.Copy(dst, src)
	if err != nil {
		log.Println(err)
		if e, ok := err.(*net.OpError); ok && e.Err == syscall.EPIPE {
			if _, ok := src.(*net.TCPConn); ok {
				src.(*net.TCPConn).CloseRead()
			} else {
				src.Close()
			}
		}
	}
	if _, ok := dst.(*net.TCPConn); ok {
		dst.(*net.TCPConn).CloseWrite()
	} else {
		dst.Close()
	}
	complete <- n
}

func (p *proxy) shutdown() {
	close(p.quit)
}

func main() {
	laddr := flag.String("l", ":8000", "Listen Address")
	saddr := flag.String("s", "", "Source Address")
	raddr := flag.String("r", "localhost:8080", "Remote Address")
	tlsAccept := flag.Bool("tls-accept", false, "Enable TLS Accept")
	tlsCert := flag.String("tls-cert", "./server.crt", "Certificate File")
	tlsKey := flag.String("tls-key", "./server.key", "Privatekey File")
	tlsConnect := flag.Bool("tls-connect", false, "Enable TLS Connect")
	flag.Parse()

	log.SetFlags(log.Ldate | log.Lmicroseconds | log.Lshortfile)
	log.Println("PID:", os.Getpid())

	config := config{
		laddr:      *laddr,
		saddr:      *saddr,
		raddr:      *raddr,
		tlsAccept:  *tlsAccept,
		tlsCert:    *tlsCert,
		tlsKey:     *tlsKey,
		tlsConnect: *tlsConnect,
	}
	proxy := NewProxy(config)
	complete := make(chan bool)
	go func() {
		defer close(complete)
		proxy.Run()
	}()
	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt)
	for {
		select {
		case s := <-sigch:
			log.Println("Recieve signal:", s)
			proxy.shutdown()
		case <-complete:
			log.Println("Good bye")
			return
		}
	}
	// does not reach
}
<|endoftext|>


<commit_before>package luar

import (
	"fmt"
	"reflect"
	"strconv"
	"sync"

	"github.com/aarzilli/golua/lua"
)

// Lua proxy objects for Go slices, maps and structs
type valueProxy struct {
	value reflect.Value
	t     reflect.Type
}

const (
	cNumberMeta    = "numberMT"
	cComplexMeta   = "complexMT"
	cStringMeta    = "stringMT"
	cSliceMeta     = "sliceMT"
	cMapMeta       = "mapMT"
	cStructMeta    = "structMT"
	cInterfaceMeta = "interfaceMT"
	cChannelMeta   = "channelMT"
)

var proxyMap = map[*valueProxy]reflect.Value{}
var proxymu = &sync.Mutex{}

// commonKind returns the kind to which v1 and v2 can be converted with the
// least information loss.
func commonKind(v1, v2 reflect.Value) reflect.Kind {
	k1 := unsizedKind(v1)
	k2 := unsizedKind(v2)
	if k1 == k2 && (k1 == reflect.Uint64 || k1 == reflect.Int64) {
		return k1
	}
	if k1 == reflect.Complex128 || k2 == reflect.Complex128 {
		return reflect.Complex128
	}
	return reflect.Float64
}

func isPointerToPrimitive(v reflect.Value) bool {
	return v.Kind() == reflect.Ptr && v.Elem().IsValid() && v.Elem().Type() != nil
}

func isPredeclaredType(t reflect.Type) bool {
	return t == reflect.TypeOf(0.0) || t == reflect.TypeOf("")
}

func isValueProxy(L *lua.State, idx int) bool {
	res := false
	if L.IsUserdata(idx) {
		L.GetMetaTable(idx)
		if !L.IsNil(-1) {
			L.GetField(-1, "luago.value")
			res = !L.IsNil(-1)
			L.Pop(1)
		}
		L.Pop(1)
	}
	return res
}

func luaToGoValue(L *lua.State, idx int) (reflect.Value, reflect.Type) {
	v := LuaToGo(L, nil, idx)
	return reflect.ValueOf(v), reflect.TypeOf(v)
}

func makeValueProxy(L *lua.State, val reflect.Value, proxyMT string) {
	// The metatable needs be set up in the Lua state before the proxy is created,
	// otherwise closing the state will fail on calling the garbage collector. Not
	// really sure why this happens though...
	L.LGetMetaTable(proxyMT)
	if L.IsNil(-1) {
		flagValue := func() {
			L.SetMetaMethod("__tostring", proxy__tostring)
			L.SetMetaMethod("__gc", proxy__gc)
			L.SetMetaMethod("__eq", proxy__eq)
			L.PushBoolean(true)
			L.SetField(-2, "luago.value")
			L.Pop(1)
		}
		switch proxyMT {
		case cNumberMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", interface__index)
			L.SetMetaMethod("__lt", number__lt)
			L.SetMetaMethod("__add", number__add)
			L.SetMetaMethod("__sub", number__sub)
			L.SetMetaMethod("__mul", number__mul)
			L.SetMetaMethod("__div", number__div)
			L.SetMetaMethod("__mod", number__mod)
			L.SetMetaMethod("__pow", number__pow)
			L.SetMetaMethod("__unm", number__unm)
			flagValue()
		case cComplexMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", complex__index)
			L.SetMetaMethod("__add", number__add)
			L.SetMetaMethod("__sub", number__sub)
			L.SetMetaMethod("__mul", number__mul)
			L.SetMetaMethod("__div", number__div)
			L.SetMetaMethod("__pow", number__pow)
			L.SetMetaMethod("__unm", number__unm)
			flagValue()
		case cStringMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", string__index)
			L.SetMetaMethod("__len", string__len)
			L.SetMetaMethod("__lt", string__lt)
			L.SetMetaMethod("__concat", string__concat)
			L.SetMetaMethod("__ipairs", string__ipairs)
			L.SetMetaMethod("__pairs", string__ipairs)
			flagValue()
		case cSliceMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", slice__index)
			L.SetMetaMethod("__newindex", slice__newindex)
			L.SetMetaMethod("__len", slicemap__len)
			L.SetMetaMethod("__ipairs", slice__ipairs)
			L.SetMetaMethod("__pairs", slice__ipairs)
			flagValue()
		case cMapMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", map__index)
			L.SetMetaMethod("__newindex", map__newindex)
			L.SetMetaMethod("__len", slicemap__len)
			L.SetMetaMethod("__ipairs", map__ipairs)
			L.SetMetaMethod("__pairs", map__pairs)
			flagValue()
		case cStructMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", struct__index)
			L.SetMetaMethod("__newindex", struct__newindex)
			flagValue()
		case cInterfaceMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", interface__index)
			flagValue()
		case cChannelMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", channel__index)
			flagValue()
		}
	}
	L.Pop(1)
	rawptr := L.NewUserdata(typeof((*valueProxy)(nil)).Size())
	ptr := (*valueProxy)(rawptr)
	ptr.value = val
	ptr.t = val.Type()
	proxymu.Lock()
	proxyMap[ptr] = val
	proxymu.Unlock()
	L.LGetMetaTable(proxyMT)
	L.SetMetaTable(-2)
}

func mustUnwrapProxy(L *lua.State, idx int) interface{} {
	if !isValueProxy(L, idx) {
		RaiseError(L, fmt.Sprintf("arg #%d is not a Go object", idx))
	}
	v, _ := valueOfProxy(L, idx)
	return v.Interface()
}

func pushGoMethod(L *lua.State, name string, st reflect.Value) {
	method := st.MethodByName(name)
	if !method.IsValid() {
		T := st.Type()
		// Could not resolve this method. Perhaps it's defined on the pointer?
		if T.Kind() != reflect.Ptr {
			if st.CanAddr() {
				// If we can get a pointer directly.
				st = st.Addr()
			} else {
				// Otherwise create and initialize one.
				VP := reflect.New(T)
				VP.Elem().Set(st)
				st = VP
			}
		}
		method = st.MethodByName(name)
		assertValid(L, method, st, name, "method")
	}
	GoToLua(L, nil, method, true)
}

// pushNumberValue pushes the number resulting from an arithmetic operation.
//
// At least one operand must be a proxy for this function to be called. See the
// main documentation for the conversion rules.
func pushNumberValue(L *lua.State, i interface{}, t1, t2 reflect.Type) {
	v := reflect.ValueOf(i)
	isComplex := unsizedKind(v) == reflect.Complex128
	mt := cNumberMeta
	if isComplex {
		mt = cComplexMeta
	}
	if t1 == t2 || isPredeclaredType(t2) {
		makeValueProxy(L, v.Convert(t1), mt)
	} else if isPredeclaredType(t1) {
		makeValueProxy(L, v.Convert(t2), mt)
	} else if isComplex {
		complexType := reflect.TypeOf(0i)
		makeValueProxy(L, v.Convert(complexType), cComplexMeta)
	} else {
		L.PushNumber(valueToNumber(L, v))
	}
}

// Shorthand for kind-switches.
func unsizedKind(v reflect.Value) reflect.Kind {
	switch v.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return reflect.Int64
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		return reflect.Uint64
	case reflect.Float64, reflect.Float32:
		return reflect.Float64
	case reflect.Complex128, reflect.Complex64:
		return reflect.Complex128
	}
	return v.Kind()
}

func valueOfProxy(L *lua.State, idx int) (reflect.Value, reflect.Type) {
	vp := (*valueProxy)(L.ToUserdata(idx))
	return vp.value, vp.t
}

func valueToComplex(L *lua.State, v reflect.Value) complex128 {
	if unsizedKind(v) == reflect.Complex128 {
		return v.Complex()
	}
	return complex(valueToNumber(L, v), 0)
}

func valueToNumber(L *lua.State, v reflect.Value) float64 {
	switch unsizedKind(v) {
	case reflect.Int64:
		return float64(v.Int())
	case reflect.Uint64:
		return float64(v.Uint())
	case reflect.Float64:
		return v.Float()
	case reflect.String:
		if f, err := strconv.ParseFloat(v.String(), 64); err == nil {
			return f
		}
		RaiseError(L, "cannot convert to number")
	}
	RaiseError(L, "cannot convert to number")
	return 0
}

func valueToString(L *lua.State, v reflect.Value) string {
	switch unsizedKind(v) {
	case reflect.Uint64, reflect.Int64, reflect.Float64:
		return fmt.Sprintf("%v", valueToNumber(L, v))
	case reflect.String:
		return v.String()
	}

	RaiseError(L, "cannot convert to string")
	return ""
}
<commit_msg>Remove spurious RaiseError in valueToNumber<commit_after>package luar

import (
	"fmt"
	"reflect"
	"strconv"
	"sync"

	"github.com/aarzilli/golua/lua"
)

// Lua proxy objects for Go slices, maps and structs
type valueProxy struct {
	value reflect.Value
	t     reflect.Type
}

const (
	cNumberMeta    = "numberMT"
	cComplexMeta   = "complexMT"
	cStringMeta    = "stringMT"
	cSliceMeta     = "sliceMT"
	cMapMeta       = "mapMT"
	cStructMeta    = "structMT"
	cInterfaceMeta = "interfaceMT"
	cChannelMeta   = "channelMT"
)

var proxyMap = map[*valueProxy]reflect.Value{}
var proxymu = &sync.Mutex{}

// commonKind returns the kind to which v1 and v2 can be converted with the
// least information loss.
func commonKind(v1, v2 reflect.Value) reflect.Kind {
	k1 := unsizedKind(v1)
	k2 := unsizedKind(v2)
	if k1 == k2 && (k1 == reflect.Uint64 || k1 == reflect.Int64) {
		return k1
	}
	if k1 == reflect.Complex128 || k2 == reflect.Complex128 {
		return reflect.Complex128
	}
	return reflect.Float64
}

func isPointerToPrimitive(v reflect.Value) bool {
	return v.Kind() == reflect.Ptr && v.Elem().IsValid() && v.Elem().Type() != nil
}

func isPredeclaredType(t reflect.Type) bool {
	return t == reflect.TypeOf(0.0) || t == reflect.TypeOf("")
}

func isValueProxy(L *lua.State, idx int) bool {
	res := false
	if L.IsUserdata(idx) {
		L.GetMetaTable(idx)
		if !L.IsNil(-1) {
			L.GetField(-1, "luago.value")
			res = !L.IsNil(-1)
			L.Pop(1)
		}
		L.Pop(1)
	}
	return res
}

func luaToGoValue(L *lua.State, idx int) (reflect.Value, reflect.Type) {
	v := LuaToGo(L, nil, idx)
	return reflect.ValueOf(v), reflect.TypeOf(v)
}

func makeValueProxy(L *lua.State, val reflect.Value, proxyMT string) {
	// The metatable needs be set up in the Lua state before the proxy is created,
	// otherwise closing the state will fail on calling the garbage collector. Not
	// really sure why this happens though...
	L.LGetMetaTable(proxyMT)
	if L.IsNil(-1) {
		flagValue := func() {
			L.SetMetaMethod("__tostring", proxy__tostring)
			L.SetMetaMethod("__gc", proxy__gc)
			L.SetMetaMethod("__eq", proxy__eq)
			L.PushBoolean(true)
			L.SetField(-2, "luago.value")
			L.Pop(1)
		}
		switch proxyMT {
		case cNumberMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", interface__index)
			L.SetMetaMethod("__lt", number__lt)
			L.SetMetaMethod("__add", number__add)
			L.SetMetaMethod("__sub", number__sub)
			L.SetMetaMethod("__mul", number__mul)
			L.SetMetaMethod("__div", number__div)
			L.SetMetaMethod("__mod", number__mod)
			L.SetMetaMethod("__pow", number__pow)
			L.SetMetaMethod("__unm", number__unm)
			flagValue()
		case cComplexMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", complex__index)
			L.SetMetaMethod("__add", number__add)
			L.SetMetaMethod("__sub", number__sub)
			L.SetMetaMethod("__mul", number__mul)
			L.SetMetaMethod("__div", number__div)
			L.SetMetaMethod("__pow", number__pow)
			L.SetMetaMethod("__unm", number__unm)
			flagValue()
		case cStringMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", string__index)
			L.SetMetaMethod("__len", string__len)
			L.SetMetaMethod("__lt", string__lt)
			L.SetMetaMethod("__concat", string__concat)
			L.SetMetaMethod("__ipairs", string__ipairs)
			L.SetMetaMethod("__pairs", string__ipairs)
			flagValue()
		case cSliceMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", slice__index)
			L.SetMetaMethod("__newindex", slice__newindex)
			L.SetMetaMethod("__len", slicemap__len)
			L.SetMetaMethod("__ipairs", slice__ipairs)
			L.SetMetaMethod("__pairs", slice__ipairs)
			flagValue()
		case cMapMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", map__index)
			L.SetMetaMethod("__newindex", map__newindex)
			L.SetMetaMethod("__len", slicemap__len)
			L.SetMetaMethod("__ipairs", map__ipairs)
			L.SetMetaMethod("__pairs", map__pairs)
			flagValue()
		case cStructMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", struct__index)
			L.SetMetaMethod("__newindex", struct__newindex)
			flagValue()
		case cInterfaceMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", interface__index)
			flagValue()
		case cChannelMeta:
			L.NewMetaTable(proxyMT)
			L.SetMetaMethod("__index", channel__index)
			flagValue()
		}
	}
	L.Pop(1)
	rawptr := L.NewUserdata(typeof((*valueProxy)(nil)).Size())
	ptr := (*valueProxy)(rawptr)
	ptr.value = val
	ptr.t = val.Type()
	proxymu.Lock()
	proxyMap[ptr] = val
	proxymu.Unlock()
	L.LGetMetaTable(proxyMT)
	L.SetMetaTable(-2)
}

func mustUnwrapProxy(L *lua.State, idx int) interface{} {
	if !isValueProxy(L, idx) {
		RaiseError(L, fmt.Sprintf("arg #%d is not a Go object", idx))
	}
	v, _ := valueOfProxy(L, idx)
	return v.Interface()
}

func pushGoMethod(L *lua.State, name string, st reflect.Value) {
	method := st.MethodByName(name)
	if !method.IsValid() {
		T := st.Type()
		// Could not resolve this method. Perhaps it's defined on the pointer?
		if T.Kind() != reflect.Ptr {
			if st.CanAddr() {
				// If we can get a pointer directly.
				st = st.Addr()
			} else {
				// Otherwise create and initialize one.
				VP := reflect.New(T)
				VP.Elem().Set(st)
				st = VP
			}
		}
		method = st.MethodByName(name)
		assertValid(L, method, st, name, "method")
	}
	GoToLua(L, nil, method, true)
}

// pushNumberValue pushes the number resulting from an arithmetic operation.
//
// At least one operand must be a proxy for this function to be called. See the
// main documentation for the conversion rules.
func pushNumberValue(L *lua.State, i interface{}, t1, t2 reflect.Type) {
	v := reflect.ValueOf(i)
	isComplex := unsizedKind(v) == reflect.Complex128
	mt := cNumberMeta
	if isComplex {
		mt = cComplexMeta
	}
	if t1 == t2 || isPredeclaredType(t2) {
		makeValueProxy(L, v.Convert(t1), mt)
	} else if isPredeclaredType(t1) {
		makeValueProxy(L, v.Convert(t2), mt)
	} else if isComplex {
		complexType := reflect.TypeOf(0i)
		makeValueProxy(L, v.Convert(complexType), cComplexMeta)
	} else {
		L.PushNumber(valueToNumber(L, v))
	}
}

// Shorthand for kind-switches.
func unsizedKind(v reflect.Value) reflect.Kind {
	switch v.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return reflect.Int64
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		return reflect.Uint64
	case reflect.Float64, reflect.Float32:
		return reflect.Float64
	case reflect.Complex128, reflect.Complex64:
		return reflect.Complex128
	}
	return v.Kind()
}

func valueOfProxy(L *lua.State, idx int) (reflect.Value, reflect.Type) {
	vp := (*valueProxy)(L.ToUserdata(idx))
	return vp.value, vp.t
}

func valueToComplex(L *lua.State, v reflect.Value) complex128 {
	if unsizedKind(v) == reflect.Complex128 {
		return v.Complex()
	}
	return complex(valueToNumber(L, v), 0)
}

func valueToNumber(L *lua.State, v reflect.Value) float64 {
	switch unsizedKind(v) {
	case reflect.Int64:
		return float64(v.Int())
	case reflect.Uint64:
		return float64(v.Uint())
	case reflect.Float64:
		return v.Float()
	case reflect.String:
		if f, err := strconv.ParseFloat(v.String(), 64); err == nil {
			return f
		}
	}
	RaiseError(L, "cannot convert to number")
	return 0
}

func valueToString(L *lua.State, v reflect.Value) string {
	switch unsizedKind(v) {
	case reflect.Uint64, reflect.Int64, reflect.Float64:
		return fmt.Sprintf("%v", valueToNumber(L, v))
	case reflect.String:
		return v.String()
	}

	RaiseError(L, "cannot convert to string")
	return ""
}
<|endoftext|>


<commit_before>package tools

import (
	"io"
	"log"
	"os"
	"os/signal"
	"syscall"
    "sync"
	"time"
)

type AppendStruct struct {
	file   *os.File
	c      chan os.Signal
	cClose chan bool
	fn     string
	err    error
    pid    int
    lock   sync.Mutex
}

func openFile(fn string) (*os.File, error) {
	return os.OpenFile(fn, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0666)
}

func (as *AppendStruct) wait() {
	for {
		select {
		case <-as.c:
			io.WriteString(os.Stderr, "we got an signal, restart at "+TimeStr())
			f, err := openFile(as.fn)
			if err != nil {
				as.err = err
				log.Println("open file:", err)
				time.Sleep(time.Second)
				continue
			}
            as.lock.Lock()
            if as.file != nil {
                as.file.Close()
            }
			as.file = f
            as.lock.Unlock()
		case <-as.cClose:
			signal.Reset(syscall.SIGUSR1)
			return
		}
	}
}

func NewAppender(fn string) (*AppendStruct, error) {
	var as AppendStruct
	f, err := openFile(fn)
	if err != nil {
		return nil, err
	}
	as.file = f
	as.fn = fn
	as.c = make(chan os.Signal)
	as.cClose = make(chan bool)
    as.pid = os.Getpid()
    signal.Notify(as.c, syscall.SIGUSR1)
	go as.wait()
	return &as, nil
}

func (as *AppendStruct) Close() {
    as.lock.Lock()
    defer as.lock.Unlock()
	as.cClose <- true
	as.file.Close()
    as.file = nil
	close(as.c)
	close(as.cClose)
}

func (as *AppendStruct) Write(bt []byte) (int, error) {
	if as.err != nil {
		return 0, as.err
	}
    as.lock.Lock()
    defer as.lock.Unlock()
	n, err := as.file.Write(bt)
	if err != nil {
		as.err = err
	}
	return n, err
}

func Daemon() {
	os.Stdout.Close()
	os.Stdin.Close()
	os.Stdout = nil
	os.Stdin = nil
}

func InitLog(fn, prefix string) *log.Logger {
	as, err := NewAppender(fn)
	if err != nil {
		return nil
	}
	if prefix == "" {
		prefix = "DefAppendLogger "
	}
	if prefix[len(prefix)-1] != ' ' {
		prefix = prefix + " "
	}
	return log.New(as, prefix, log.LstdFlags|log.Lshortfile)
}
<commit_msg>fix sig error<commit_after>package tools

import (
	"io"
	"log"
	"os"
	"os/signal"
	"syscall"
    "sync"
	"time"
)

type AppendStruct struct {
	file   *os.File
	cSig      chan os.Signal
	cClose chan bool
	fn     string
	err    error
    pid    int
    lock   sync.Mutex
}

func openFile(fn string) (*os.File, error) {
	return os.OpenFile(fn, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0666)
}

func (as *AppendStruct) wait() {
	for {
		select {
		case _, ok := <-as.cSig:
            if ok == false {
                return
            }
			io.WriteString(os.Stderr, "we got an signal, restart at "+TimeStr())
			f, err := openFile(as.fn)
			if err != nil {
				as.err = err
				log.Println("open file:", err)
				time.Sleep(time.Second)
				continue
			}
            as.lock.Lock()
            if as.file != nil {
                as.file.Close()
            }
			as.file = f
            as.lock.Unlock()
		case <-as.cClose:
			signal.Reset(syscall.SIGUSR1)
		}
	}
}

func NewAppender(fn string) (*AppendStruct, error) {
	var as AppendStruct
	f, err := openFile(fn)
	if err != nil {
		return nil, err
	}
	as.file = f
	as.fn = fn
	as.cSig = make(chan os.Signal)
	as.cClose = make(chan bool)
    as.pid = os.Getpid()
    signal.Notify(as.cSig, syscall.SIGUSR1)
	go as.wait()
	return &as, nil
}

func (as *AppendStruct) Close() {
    as.lock.Lock()
    defer as.lock.Unlock()
	as.cClose <- true
    signal.Stop(as.cSig)
	close(as.cSig)
    as.file.Close()
    as.file = nil
	close(as.cClose)
}

func (as *AppendStruct) Write(bt []byte) (int, error) {
	if as.err != nil {
		return 0, as.err
	}
    as.lock.Lock()
    defer as.lock.Unlock()
	n, err := as.file.Write(bt)
	if err != nil {
		as.err = err
	}
	return n, err
}

func Daemon() {
	os.Stdout.Close()
	os.Stdin.Close()
	os.Stdout = nil
	os.Stdin = nil
}

func InitLog(fn, prefix string) *log.Logger {
	as, err := NewAppender(fn)
	if err != nil {
		return nil
	}
	if prefix == "" {
		prefix = "DefAppendLogger "
	}
	if prefix[len(prefix)-1] != ' ' {
		prefix = prefix + " "
	}
	return log.New(as, prefix, log.LstdFlags|log.Lshortfile)
}
<|endoftext|>


<commit_before>// Copyright 2014 Joseph Hager. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package engi

import (
	"math"
	"path"
	"path/filepath"

	"github.com/ajhager/webgl"
)

type Resource struct {
	kind string
	name string
	url  string
}

type Loader struct {
	resources []Resource
	loaded    map[Resource]struct{}
	images    map[string]*Texture
	jsons     map[string]string
	sounds    map[string]*Sound
}

func NewLoader() *Loader {
	return &Loader{
		resources: make([]Resource, 1),
		loaded:    make(map[Resource]struct{}),
		images:    make(map[string]*Texture),
		jsons:     make(map[string]string),
		sounds:    make(map[string]*Sound),
	}
}

func (l *Loader) Add(name, url string) {
	kind := path.Ext(url)[1:]
	url = filepath.FromSlash(url)
	l.resources = append(l.resources, Resource{kind, name, url})
}

func (l *Loader) Image(name string) *Texture {
	return l.images[name]
}

func (l *Loader) Json(name string) string {
	return l.jsons[name]
}

func (l *Loader) Sound(name string) *Sound {
	return l.sounds[name]
}

func (l *Loader) Replace(name, url string) {
	kind := path.Ext(url)[1:]
	url = filepath.FromSlash(url)
	r := Resource{kind, name, url}
	switch kind {
	case "wav", "flac":
		if old, ok := l.sounds[name]; ok {
			old.Delete()
		}
		if data, err := loadSound(r); err == nil {
			l.sounds[r.name] = data
			l.loaded[r] = struct{}{}
		} else {
			panic(err)
		}
	}
}

func (l *Loader) Load(onFinish func()) {
	for _, r := range l.resources {
		if _, loaded := l.loaded[r]; loaded {
			// don't load stuff twice
			continue
		}
		switch r.kind {
		case "png":
			data, err := loadImage(r)
			if err == nil {
				l.images[r.name] = NewTexture(data)
				l.loaded[r] = struct{}{}
			}
		case "json":
			data, err := loadJson(r)
			if err == nil {
				l.jsons[r.name] = data
				l.loaded[r] = struct{}{}
			}
		case "wav", "flac":
			data, err := loadSound(r)
			if err == nil {
				l.sounds[r.name] = data
				l.loaded[r] = struct{}{}
			} else {
				panic(err)
			}
		}
	}
	onFinish()
}

type Image interface {
	Data() interface{}
	Width() int
	Height() int
}

func LoadShader(vertSrc, fragSrc string) *webgl.Program {
	vertShader := gl.CreateShader(gl.VERTEX_SHADER)
	gl.ShaderSource(vertShader, vertSrc)
	gl.CompileShader(vertShader)
	defer gl.DeleteShader(vertShader)

	fragShader := gl.CreateShader(gl.FRAGMENT_SHADER)
	gl.ShaderSource(fragShader, fragSrc)
	gl.CompileShader(fragShader)
	defer gl.DeleteShader(fragShader)

	program := gl.CreateProgram()
	gl.AttachShader(program, vertShader)
	gl.AttachShader(program, fragShader)
	gl.LinkProgram(program)

	return program
}

type Region struct {
	texture       *Texture
	u, v          float32
	u2, v2        float32
	width, height float32
}

func NewRegion(texture *Texture, x, y, w, h int) *Region {
	invTexWidth := 1.0 / float32(texture.Width())
	invTexHeight := 1.0 / float32(texture.Height())

	u := float32(x) * invTexWidth
	v := float32(y) * invTexHeight
	u2 := float32(x+w) * invTexWidth
	v2 := float32(y+h) * invTexHeight
	width := float32(math.Abs(float64(w)))
	height := float32(math.Abs(float64(h)))

	return &Region{texture, u, v, u2, v2, width, height}
}

func (r *Region) Width() float32 {
	return float32(r.width)
}

func (r *Region) Height() float32 {
	return float32(r.height)
}

func (r *Region) Texture() *webgl.Texture {
	return r.texture.id
}

func (r *Region) View() (float32, float32, float32, float32) {
	return r.u, r.v, r.u2, r.v2
}

type Texture struct {
	id     *webgl.Texture
	width  int
	height int
}

func NewTexture(img Image) *Texture {
	id := gl.CreateTexture()

	gl.BindTexture(gl.TEXTURE_2D, id)

	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE)
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE)
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR)
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST)

	if img.Data() == nil {
		panic("Texture image data is nil.")
	}

	gl.TexImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, img.Data())

	return &Texture{id, img.Width(), img.Height()}
}

// Width returns the width of the texture.
func (t *Texture) Width() float32 {
	return float32(t.width)
}

// Height returns the height of the texture.
func (t *Texture) Height() float32 {
	return float32(t.height)
}

func (t *Texture) Texture() *webgl.Texture {
	return t.id
}

func (r *Texture) View() (float32, float32, float32, float32) {
	return 0.0, 0.0, 1.0, 1.0
}

type Point struct {
	X, Y float32
}

func (p *Point) Set(x, y float32) {
	p.X = x
	p.Y = y
}

func (p *Point) SetTo(v float32) {
	p.X = v
	p.Y = v
}

type Sprite struct {
	Position *Point
	Scale    *Point
	Anchor   *Point
	Rotation float32
	Color    uint32
	Alpha    float32
	Region   *Region
}

func NewSprite(region *Region, x, y float32) *Sprite {
	return &Sprite{
		Position: &Point{x, y},
		Scale:    &Point{1, 1},
		Anchor:   &Point{0, 0},
		Rotation: 0,
		Color:    0xffffff,
		Alpha:    1,
		Region:   region,
	}
}

func (s *Sprite) Render(batch *Batch) {
	batch.Draw(s.Region, s.Position.X, s.Position.Y, s.Anchor.X, s.Anchor.Y, s.Scale.X, s.Scale.Y, s.Rotation, s.Color, s.Alpha)
}

var batchVert = `
attribute vec2 in_Position;
attribute vec4 in_Color;
attribute vec2 in_TexCoords;

uniform vec2 uf_Projection;

varying vec4 var_Color;
varying vec2 var_TexCoords;

const vec2 center = vec2(-1.0, 1.0);

void main() {
  var_Color = in_Color;
  var_TexCoords = in_TexCoords;
	gl_Position = vec4(in_Position.x / uf_Projection.x + center.x,
										 in_Position.y / -uf_Projection.y + center.y,
										 0.0, 1.0);
}`

var batchFrag = `
#ifdef GL_ES
#define LOWP lowp
precision mediump float;
#else
#define LOWP
#endif

varying vec4 var_Color;
varying vec2 var_TexCoords;

uniform sampler2D uf_Texture;

void main (void) {
  gl_FragColor = var_Color * texture2D(uf_Texture, var_TexCoords);
}`
<commit_msg>panic on asset loading errors<commit_after>// Copyright 2014 Joseph Hager. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package engi

import (
	"math"
	"path"
	"path/filepath"

	"github.com/ajhager/webgl"
)

type Resource struct {
	kind string
	name string
	url  string
}

type Loader struct {
	resources []Resource
	loaded    map[Resource]struct{}
	images    map[string]*Texture
	jsons     map[string]string
	sounds    map[string]*Sound
}

func NewLoader() *Loader {
	return &Loader{
		resources: make([]Resource, 1),
		loaded:    make(map[Resource]struct{}),
		images:    make(map[string]*Texture),
		jsons:     make(map[string]string),
		sounds:    make(map[string]*Sound),
	}
}

func (l *Loader) Add(name, url string) {
	kind := path.Ext(url)[1:]
	url = filepath.FromSlash(url)
	l.resources = append(l.resources, Resource{kind, name, url})
}

func (l *Loader) Image(name string) *Texture {
	return l.images[name]
}

func (l *Loader) Json(name string) string {
	return l.jsons[name]
}

func (l *Loader) Sound(name string) *Sound {
	return l.sounds[name]
}

func (l *Loader) Replace(name, url string) {
	kind := path.Ext(url)[1:]
	url = filepath.FromSlash(url)
	r := Resource{kind, name, url}
	switch kind {
	case "wav", "flac":
		if old, ok := l.sounds[name]; ok {
			old.Delete()
		}
		data, err := loadSound(r)
		fatalErr(err)
		l.sounds[r.name] = data
		l.loaded[r] = struct{}{}
	}
}

func (l *Loader) Load(onFinish func()) {
	for _, r := range l.resources {
		if _, loaded := l.loaded[r]; loaded {
			// don't load stuff twice
			continue
		}
		switch r.kind {
		case "png":
			data, err := loadImage(r)
			fatalErr(err)
			l.images[r.name] = NewTexture(data)
			l.loaded[r] = struct{}{}
		case "json":
			data, err := loadJson(r)
			fatalErr(err)
			l.jsons[r.name] = data
			l.loaded[r] = struct{}{}
		case "wav", "flac":
			data, err := loadSound(r)
			fatalErr(err)
			l.sounds[r.name] = data
			l.loaded[r] = struct{}{}
		}
	}
	onFinish()
}

type Image interface {
	Data() interface{}
	Width() int
	Height() int
}

func LoadShader(vertSrc, fragSrc string) *webgl.Program {
	vertShader := gl.CreateShader(gl.VERTEX_SHADER)
	gl.ShaderSource(vertShader, vertSrc)
	gl.CompileShader(vertShader)
	defer gl.DeleteShader(vertShader)

	fragShader := gl.CreateShader(gl.FRAGMENT_SHADER)
	gl.ShaderSource(fragShader, fragSrc)
	gl.CompileShader(fragShader)
	defer gl.DeleteShader(fragShader)

	program := gl.CreateProgram()
	gl.AttachShader(program, vertShader)
	gl.AttachShader(program, fragShader)
	gl.LinkProgram(program)

	return program
}

type Region struct {
	texture       *Texture
	u, v          float32
	u2, v2        float32
	width, height float32
}

func NewRegion(texture *Texture, x, y, w, h int) *Region {
	invTexWidth := 1.0 / float32(texture.Width())
	invTexHeight := 1.0 / float32(texture.Height())

	u := float32(x) * invTexWidth
	v := float32(y) * invTexHeight
	u2 := float32(x+w) * invTexWidth
	v2 := float32(y+h) * invTexHeight
	width := float32(math.Abs(float64(w)))
	height := float32(math.Abs(float64(h)))

	return &Region{texture, u, v, u2, v2, width, height}
}

func (r *Region) Width() float32 {
	return float32(r.width)
}

func (r *Region) Height() float32 {
	return float32(r.height)
}

func (r *Region) Texture() *webgl.Texture {
	return r.texture.id
}

func (r *Region) View() (float32, float32, float32, float32) {
	return r.u, r.v, r.u2, r.v2
}

type Texture struct {
	id     *webgl.Texture
	width  int
	height int
}

func NewTexture(img Image) *Texture {
	id := gl.CreateTexture()

	gl.BindTexture(gl.TEXTURE_2D, id)

	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE)
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE)
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR)
	gl.TexParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST)

	if img.Data() == nil {
		panic("Texture image data is nil.")
	}

	gl.TexImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, img.Data())

	return &Texture{id, img.Width(), img.Height()}
}

// Width returns the width of the texture.
func (t *Texture) Width() float32 {
	return float32(t.width)
}

// Height returns the height of the texture.
func (t *Texture) Height() float32 {
	return float32(t.height)
}

func (t *Texture) Texture() *webgl.Texture {
	return t.id
}

func (r *Texture) View() (float32, float32, float32, float32) {
	return 0.0, 0.0, 1.0, 1.0
}

type Point struct {
	X, Y float32
}

func (p *Point) Set(x, y float32) {
	p.X = x
	p.Y = y
}

func (p *Point) SetTo(v float32) {
	p.X = v
	p.Y = v
}

type Sprite struct {
	Position *Point
	Scale    *Point
	Anchor   *Point
	Rotation float32
	Color    uint32
	Alpha    float32
	Region   *Region
}

func NewSprite(region *Region, x, y float32) *Sprite {
	return &Sprite{
		Position: &Point{x, y},
		Scale:    &Point{1, 1},
		Anchor:   &Point{0, 0},
		Rotation: 0,
		Color:    0xffffff,
		Alpha:    1,
		Region:   region,
	}
}

func (s *Sprite) Render(batch *Batch) {
	batch.Draw(s.Region, s.Position.X, s.Position.Y, s.Anchor.X, s.Anchor.Y, s.Scale.X, s.Scale.Y, s.Rotation, s.Color, s.Alpha)
}

var batchVert = `
attribute vec2 in_Position;
attribute vec4 in_Color;
attribute vec2 in_TexCoords;

uniform vec2 uf_Projection;

varying vec4 var_Color;
varying vec2 var_TexCoords;

const vec2 center = vec2(-1.0, 1.0);

void main() {
  var_Color = in_Color;
  var_TexCoords = in_TexCoords;
	gl_Position = vec4(in_Position.x / uf_Projection.x + center.x,
										 in_Position.y / -uf_Projection.y + center.y,
										 0.0, 1.0);
}`

var batchFrag = `
#ifdef GL_ES
#define LOWP lowp
precision mediump float;
#else
#define LOWP
#endif

varying vec4 var_Color;
varying vec2 var_TexCoords;

uniform sampler2D uf_Texture;

void main (void) {
  gl_FragColor = var_Color * texture2D(uf_Texture, var_TexCoords);
}`
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/dchest/uniuri"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/mux"
	"github.com/mholt/binding"
	"github.com/phyber/negroni-gzip/gzip"
	"github.com/rs/cors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"
)

const cookieMaxAge = 60 * 60 * 60 * 24 * 30

var (
	pool *redis.Pool
	png  = mustReadFile("assets/beacon.png")
)

func mustReadFile(path string) []byte {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return b
}

type TrackJson struct {
	Visits  int64 `json:"visits"`
	Uniques int64 `json:"uniques"`
}

func (trackJson *TrackJson) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&trackJson.Visits:  "visits",
		&trackJson.Uniques: "uniques",
	}
}

func uid(w http.ResponseWriter, req *http.Request) string {
	cookie, err := req.Cookie("uid")
	if err != nil {
		switch err {
		case http.ErrNoCookie:
			uid := fmt.Sprintf("%s", uniuri.New())
			now := time.Now()
			new_cookie := &http.Cookie{Name: "uid", Value: uid, MaxAge: cookieMaxAge, Expires: now.Add(cookieMaxAge)}
			log.Print("Setting new cookie ", new_cookie)
			http.SetCookie(w, new_cookie)
			return uid
		default:
			log.Fatal(err)
			return ""
		}
	}
	return cookie.Value
}

func track(objectId string, uid string) {
	log.Print("Tracking ", uid, " on ", objectId)
	conn := pool.Get()
	defer conn.Close()

	// http://godoc.org/github.com/garyburd/redigo/redis#hdr-Pipelining
	conn.Send("MULTI")

	// Track the number of unique visitors in a HyperLogLog
	// http://redis.io/commands/pfadd
	conn.Send("PFADD", "hll_"+objectId, uid)

	// Track the total number of visits in a simple key (stringy)
	// http://redis.io/commands/incr
	conn.Send("INCR", "str_"+objectId)

	_, err := conn.Do("EXEC")
	if err != nil {
		log.Print(err)
		return
	}
}

func beaconHandler(w http.ResponseWriter, req *http.Request) {
	fmt.Println(req.Method, req.URL)
	query, _ := url.ParseQuery(req.URL.RawQuery)
	objectId := query.Get("id")
	if objectId != "" {
		go track(objectId, uid(w, req))
	}
	w.Header().Set("Content-Type", "image/png")
	w.Write(png)
}

func apiHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	objectId := vars["objectId"]
	conn := pool.Get()
	defer conn.Close()
	visits, _ := redis.Int64(conn.Do("GET", "str_"+objectId))
	uniques, _ := redis.Int64(conn.Do("PFCOUNT", "hll_"+objectId))
	apiResponse := TrackJson{Visits: visits, Uniques: uniques}
	js, err := json.Marshal(apiResponse)
	if err != nil {
		log.Print(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func apiWriteHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	objectId := vars["objectId"]
	trackJson := new(TrackJson)
	if binding.Bind(req, trackJson).Handle(w) {
	    return
	}
	fmt.Sprintf("%q\n", trackJson)
	conn := pool.Get()
	defer conn.Close()
	_, err := conn.Do("MSET", "uniques_"+objectId, trackJson.Uniques, "visits_"+objectId, trackJson.Visits)
	if err != nil {
		log.Print(err)

	}
}

func listenAddress() string {
	string := os.Getenv("PORT")
	if string == "" {
		return ":8080"
	} else {
		return ":" + string
	}
}

func redisConfig() (string, string) {
	redis_provider := os.Getenv("REDIS_PROVIDER")
	if redis_provider == "" {
		redis_provider = "OPENREDIS_URL"
	}
	string := os.Getenv(redis_provider)
	if string != "" {
		url, err := url.Parse(string)
		password := ""
		if err != nil {
			panic(err)
		}
		if url.User != nil {
			password, _ = url.User.Password()
		}
		return url.Host, password
	} else {
		return "127.0.0.1:6379", ""

	}
}

func newPool(server, password string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", server)
			if err != nil {
				return nil, err
			}
			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					return nil, err
				}
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

func main() {
	redisServer, redisPassword := redisConfig()
	log.Print("Connecting to Redis on ", redisServer, redisPassword)
	pool = newPool(redisServer, redisPassword)

	r := mux.NewRouter()
	r.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		http.Redirect(w, req, "https://www.github.com/jelder/beacon", 302)
	})
	r.HandleFunc("/beacon.png", beaconHandler)
	r.HandleFunc("/api/{objectId}", apiHandler).Methods("GET")
	r.HandleFunc("/api/{objectId}", apiWriteHandler).Methods("POST")

	n := negroni.Classic()
	n.Use(gzip.Gzip(gzip.DefaultCompression))
	n.Use(cors.New(cors.Options{
		AllowedOrigins: []string{"*"},
	}))
	n.UseHandler(r)
	n.Run(listenAddress())
}
<commit_msg>Combine migrated and real values in response<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/dchest/uniuri"
	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/mux"
	"github.com/mholt/binding"
	"github.com/phyber/negroni-gzip/gzip"
	"github.com/rs/cors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"
)

const cookieMaxAge = 60 * 60 * 60 * 24 * 30

var (
	pool *redis.Pool
	png  = mustReadFile("assets/beacon.png")
)

func mustReadFile(path string) []byte {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return b
}

type TrackJson struct {
	Visits  int64 `json:"visits"`
	Uniques int64 `json:"uniques"`
}

func (trackJson *TrackJson) FieldMap() binding.FieldMap {
	return binding.FieldMap{
		&trackJson.Visits:  "visits",
		&trackJson.Uniques: "uniques",
	}
}

func uid(w http.ResponseWriter, req *http.Request) string {
	cookie, err := req.Cookie("uid")
	if err != nil {
		switch err {
		case http.ErrNoCookie:
			uid := fmt.Sprintf("%s", uniuri.New())
			now := time.Now()
			new_cookie := &http.Cookie{Name: "uid", Value: uid, MaxAge: cookieMaxAge, Expires: now.Add(cookieMaxAge)}
			log.Print("Setting new cookie ", new_cookie)
			http.SetCookie(w, new_cookie)
			return uid
		default:
			log.Fatal(err)
			return ""
		}
	}
	return cookie.Value
}

func track(objectId string, uid string) {
	log.Print("Tracking ", uid, " on ", objectId)
	conn := pool.Get()
	defer conn.Close()

	// http://godoc.org/github.com/garyburd/redigo/redis#hdr-Pipelining
	conn.Send("MULTI")

	// Track the number of unique visitors in a HyperLogLog
	// http://redis.io/commands/pfadd
	conn.Send("PFADD", "hll_"+objectId, uid)

	// Track the total number of visits in a simple key (stringy)
	// http://redis.io/commands/incr
	conn.Send("INCR", "str_"+objectId)

	_, err := conn.Do("EXEC")
	if err != nil {
		log.Print(err)
		return
	}
}

func beaconHandler(w http.ResponseWriter, req *http.Request) {
	fmt.Println(req.Method, req.URL)
	query, _ := url.ParseQuery(req.URL.RawQuery)
	objectId := query.Get("id")
	if objectId != "" {
		go track(objectId, uid(w, req))
	}
	w.Header().Set("Content-Type", "image/png")
	w.Write(png)
}

func apiHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	objectId := vars["objectId"]
	conn := pool.Get()
	defer conn.Close()

	uniques, err := redis.Int64(conn.Do("PFCOUNT", "hll_"+objectId))
	if err != nil {
		log.Print(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	var migrated_visits, migrated_uniques, visits int64
	mget, err := redis.Values(conn.Do("MGET", "visits_"+objectId, "uniques_"+objectId, "str_"+objectId))
	if err != nil {
		log.Print(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if _, err := redis.Scan(mget, &migrated_visits, &migrated_uniques, &visits); err != nil {
		log.Print(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	visits += migrated_visits
	uniques += migrated_uniques

	apiResponse := TrackJson{Visits: visits, Uniques: uniques}
	js, _ := json.Marshal(apiResponse)
	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func apiWriteHandler(w http.ResponseWriter, req *http.Request) {
	vars := mux.Vars(req)
	objectId := vars["objectId"]
	trackJson := new(TrackJson)
	if binding.Bind(req, trackJson).Handle(w) {
		return
	}
	fmt.Sprintf("%q\n", trackJson)
	conn := pool.Get()
	defer conn.Close()
	_, err := conn.Do("MSET", "uniques_"+objectId, trackJson.Uniques, "visits_"+objectId, trackJson.Visits)
	if err != nil {
		log.Print(err)

	}
}

func listenAddress() string {
	string := os.Getenv("PORT")
	if string == "" {
		return ":8080"
	} else {
		return ":" + string
	}
}

func redisConfig() (string, string) {
	redis_provider := os.Getenv("REDIS_PROVIDER")
	if redis_provider == "" {
		redis_provider = "OPENREDIS_URL"
	}
	string := os.Getenv(redis_provider)
	if string != "" {
		url, err := url.Parse(string)
		password := ""
		if err != nil {
			panic(err)
		}
		if url.User != nil {
			password, _ = url.User.Password()
		}
		return url.Host, password
	} else {
		return "127.0.0.1:6379", ""

	}
}

func newPool(server, password string) *redis.Pool {
	return &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", server)
			if err != nil {
				return nil, err
			}
			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					return nil, err
				}
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

func main() {
	redisServer, redisPassword := redisConfig()
	log.Print("Connecting to Redis on ", redisServer, redisPassword)
	pool = newPool(redisServer, redisPassword)

	r := mux.NewRouter()
	r.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		http.Redirect(w, req, "https://www.github.com/jelder/beacon", 302)
	})
	r.HandleFunc("/beacon.png", beaconHandler)
	r.HandleFunc("/api/{objectId}", apiHandler).Methods("GET")
	r.HandleFunc("/api/{objectId}", apiWriteHandler).Methods("POST")

	n := negroni.Classic()
	n.Use(gzip.Gzip(gzip.DefaultCompression))
	n.Use(cors.New(cors.Options{
		AllowedOrigins: []string{"*"},
	}))
	n.UseHandler(r)
	n.Run(listenAddress())
}
<|endoftext|>


<commit_before>package main

import (
	"time"
	"unicode/utf8"

	"github.com/256dpi/fire"
	"github.com/256dpi/fire/coal"
	"github.com/256dpi/fire/flame"

	"github.com/globalsign/mgo/bson"
)

var catalog = coal.NewCatalog(&Item{})

var indexer = coal.NewIndexer()

func init() {
	// add flame indexes
	flame.AddApplicationIndexes(indexer)
	flame.AddUserIndexes(indexer)
	flame.AddAccessTokenIndexes(indexer, true)
	flame.AddRefreshTokenIndexes(indexer, true)

	// add item index
	indexer.Add(&Item{}, false, 0, "Name")

	// add background delete index
	indexer.Add(&Item{}, false, time.Second, "Deleted")
}

// EnsureIndexes will ensure that the required indexes exist.
func EnsureIndexes(store *coal.Store) error {
	// ensure model indexes
	err := indexer.Ensure(store)
	if err != nil {
		return err
	}

	return nil
}

// Item represents a general item.
type Item struct {
	coal.Base `json:"-" bson:",inline" coal:"items"`
	Name      string     `json:"name"`
	State     bool       `json:"state"`
	Count     int        `json:"count"`
	Created   time.Time  `json:"created-at" bson:"created_at"`
	Deleted   *time.Time `json:"deleted-at" bson:"deleted_at"`
}

// Validate implements the fire.ValidatableModel interface.
func (i *Item) Validate() error {
	// check name
	if utf8.RuneCountInString(i.Name) < 1 {
		return fire.E("missing name")
	}

	// check created at
	if i.Created.IsZero() {
		return fire.E("missing timestamp")
	}

	return nil
}

// SoftDeleteField implements the fire.SoftDeletableModel interface.
func (i *Item) SoftDeleteField() string {
	return "Deleted"
}

type count struct {
	Item bson.ObjectId `bson:"item_id"`
}
<commit_msg>added axe indexes<commit_after>package main

import (
	"time"
	"unicode/utf8"

	"github.com/256dpi/fire"
	"github.com/256dpi/fire/axe"
	"github.com/256dpi/fire/coal"
	"github.com/256dpi/fire/flame"

	"github.com/globalsign/mgo/bson"
)

var catalog = coal.NewCatalog(&Item{})

var indexer = coal.NewIndexer()

func init() {
	// add flame indexes
	flame.AddApplicationIndexes(indexer)
	flame.AddUserIndexes(indexer)
	flame.AddAccessTokenIndexes(indexer, true)
	flame.AddRefreshTokenIndexes(indexer, true)

	// add axe indexes
	axe.AddJobIndexes(indexer, time.Minute)

	// add item index
	indexer.Add(&Item{}, false, 0, "Name")

	// add background delete index
	indexer.Add(&Item{}, false, time.Second, "Deleted")
}

// EnsureIndexes will ensure that the required indexes exist.
func EnsureIndexes(store *coal.Store) error {
	// ensure model indexes
	err := indexer.Ensure(store)
	if err != nil {
		return err
	}

	return nil
}

// Item represents a general item.
type Item struct {
	coal.Base `json:"-" bson:",inline" coal:"items"`
	Name      string     `json:"name"`
	State     bool       `json:"state"`
	Count     int        `json:"count"`
	Created   time.Time  `json:"created-at" bson:"created_at"`
	Deleted   *time.Time `json:"deleted-at" bson:"deleted_at"`
}

// Validate implements the fire.ValidatableModel interface.
func (i *Item) Validate() error {
	// check name
	if utf8.RuneCountInString(i.Name) < 1 {
		return fire.E("missing name")
	}

	// check created at
	if i.Created.IsZero() {
		return fire.E("missing timestamp")
	}

	return nil
}

// SoftDeleteField implements the fire.SoftDeletableModel interface.
func (i *Item) SoftDeleteField() string {
	return "Deleted"
}

type count struct {
	Item bson.ObjectId `bson:"item_id"`
}
<|endoftext|>