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

package metrics

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"time"

	onet "github.com/Jigsaw-Code/outline-ss-server/net"
	geoip2 "github.com/oschwald/geoip2-golang"
	"github.com/prometheus/client_golang/prometheus"
)

// ShadowsocksMetrics registers metrics for the Shadowsocks service.
type ShadowsocksMetrics interface {
	SetBuildInfo(version string)

	GetLocation(net.Addr) (string, error)

	SetNumAccessKeys(numKeys int, numPorts int)

	// TCP metrics
	AddOpenTCPConnection(clientLocation string)
	AddClosedTCPConnection(clientLocation, accessKey, status string, data ProxyMetrics, timeToCipher, duration time.Duration)
	AddTCPProbe(clientLocation, status, drainResult string, port int, data ProxyMetrics)

	// UDP metrics
	AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration)
	AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int)
	AddUDPNatEntry()
	RemoveUDPNatEntry()
}

type shadowsocksMetrics struct {
	ipCountryDB *geoip2.Reader

	buildInfo      *prometheus.GaugeVec
	accessKeys     prometheus.Gauge
	ports          prometheus.Gauge
	dataBytes      *prometheus.CounterVec
	timeToCipherMs *prometheus.HistogramVec
	// TODO: Add time to first byte.

	tcpProbes               *prometheus.HistogramVec
	tcpOpenConnections      *prometheus.CounterVec
	tcpClosedConnections    *prometheus.CounterVec
	tcpConnectionDurationMs *prometheus.HistogramVec

	udpAddedNatEntries   prometheus.Counter
	udpRemovedNatEntries prometheus.Counter
}

func newShadowsocksMetrics(ipCountryDB *geoip2.Reader) *shadowsocksMetrics {
	return &shadowsocksMetrics{
		ipCountryDB: ipCountryDB,
		buildInfo: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Namespace: "shadowsocks",
			Name:      "build_info",
			Help:      "Information on the outline-ss-server build",
		}, []string{"version"}),
		accessKeys: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "shadowsocks",
			Name:      "keys",
			Help:      "Count of access keys",
		}),
		ports: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "shadowsocks",
			Name:      "ports",
			Help:      "Count of open Shadowsocks ports",
		}),
		tcpOpenConnections: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: "shadowsocks",
			Subsystem: "tcp",
			Name:      "connections_opened",
			Help:      "Count of open TCP connections",
		}, []string{"location"}),
		tcpClosedConnections: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: "shadowsocks",
			Subsystem: "tcp",
			Name:      "connections_closed",
			Help:      "Count of closed TCP connections",
		}, []string{"location", "status", "access_key"}),
		tcpConnectionDurationMs: prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: "shadowsocks",
				Subsystem: "tcp",
				Name:      "connection_duration_ms",
				Help:      "TCP connection duration distributions.",
				Buckets: []float64{
					100,
					float64(time.Second.Milliseconds()),
					float64(time.Minute.Milliseconds()),
					float64(time.Hour.Milliseconds()),
					float64(24 * time.Hour.Milliseconds()),     // Day
					float64(7 * 24 * time.Hour.Milliseconds()), // Week
				},
			}, []string{"status"}),
		dataBytes: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: "shadowsocks",
				Name:      "data_bytes",
				Help:      "Bytes transferred by the proxy",
			}, []string{"dir", "proto", "location", "access_key"}),
		tcpProbes: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Namespace: "shadowsocks",
			Name:      "tcp_probes",
			Buckets:   []float64{0, 48, 49, 50, 51, 52, 72, 73, 90, 91, 220, 221},
			Help:      "Histogram of number of bytes from client to proxy, for detecting possible probes",
		}, []string{"location", "port", "status", "error"}),
		timeToCipherMs: prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: "shadowsocks",
				Name:      "time_to_cipher_ms",
				Help:      "Time needed to find the cipher",
				Buckets:   []float64{0.1, 1, 10, 100, 1000},
			}, []string{"proto", "found_key"}),
		udpAddedNatEntries: prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "shadowsocks",
				Subsystem: "udp",
				Name:      "nat_entries_added",
				Help:      "Entries added to the UDP NAT table",
			}),
		udpRemovedNatEntries: prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "shadowsocks",
				Subsystem: "udp",
				Name:      "nat_entries_removed",
				Help:      "Entries removed from the UDP NAT table",
			}),
	}
}

// NewPrometheusShadowsocksMetrics constructs a metrics object that uses
// `ipCountryDB` to convert IP addresses to countries, and reports all
// metrics to Prometheus via `registerer`.  `ipCountryDB` may be nil, but
// `registerer` must not be.
func NewPrometheusShadowsocksMetrics(ipCountryDB *geoip2.Reader, registerer prometheus.Registerer) ShadowsocksMetrics {
	m := newShadowsocksMetrics(ipCountryDB)
	// TODO: Is it possible to pass where to register the collectors?
	registerer.MustRegister(m.buildInfo, m.accessKeys, m.ports, m.tcpOpenConnections, m.tcpProbes, m.tcpClosedConnections, m.tcpConnectionDurationMs,
		m.dataBytes, m.timeToCipherMs, m.udpAddedNatEntries, m.udpRemovedNatEntries)
	return m
}

const (
	errParseAddr     = "XA"
	errDbLookupError = "XD"
	localLocation    = "XL"
	unknownLocation  = "ZZ"
)

func (m *shadowsocksMetrics) SetBuildInfo(version string) {
	m.buildInfo.WithLabelValues(version).Set(1)
}

func (m *shadowsocksMetrics) GetLocation(addr net.Addr) (string, error) {
	if m.ipCountryDB == nil {
		return "", nil
	}
	hostname, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return errParseAddr, errors.New("Failed to split hostname and port")
	}
	ip := net.ParseIP(hostname)
	if ip == nil {
		return errParseAddr, errors.New("Failed to parse address as IP")
	}
	if ip.IsLoopback() {
		return localLocation, nil
	}
	if !ip.IsGlobalUnicast() {
		return localLocation, nil
	}
	record, err := m.ipCountryDB.Country(ip)
	if err != nil {
		return errDbLookupError, errors.New("IP lookup failed")
	}
	if record == nil {
		return unknownLocation, errors.New("IP lookup returned nil")
	}
	if record.Country.IsoCode == "" {
		return unknownLocation, errors.New("IP Lookup has empty ISO code")
	}
	return record.Country.IsoCode, nil
}

func (m *shadowsocksMetrics) SetNumAccessKeys(numKeys int, ports int) {
	m.accessKeys.Set(float64(numKeys))
	m.ports.Set(float64(ports))
}

func (m *shadowsocksMetrics) AddOpenTCPConnection(clientLocation string) {
	m.tcpOpenConnections.WithLabelValues(clientLocation).Inc()
}

// Converts accessKey to "true" or "false"
func isFound(accessKey string) string {
	return fmt.Sprintf("%t", accessKey != "")
}

// addIfNonZero helps avoid the creation of series that are always zero.
func addIfNonZero(counter prometheus.Counter, value int64) {
	if value > 0 {
		counter.Add(float64(value))
	}
}

func (m *shadowsocksMetrics) AddClosedTCPConnection(clientLocation, accessKey, status string, data ProxyMetrics, timeToCipher, duration time.Duration) {
	m.tcpClosedConnections.WithLabelValues(clientLocation, status, accessKey).Inc()
	m.tcpConnectionDurationMs.WithLabelValues(status).Observe(duration.Seconds() * 1000)
	m.timeToCipherMs.WithLabelValues("tcp", isFound(accessKey)).Observe(timeToCipher.Seconds() * 1000)
	addIfNonZero(m.dataBytes.WithLabelValues("c>p", "tcp", clientLocation, accessKey), data.ClientProxy)
	addIfNonZero(m.dataBytes.WithLabelValues("p>t", "tcp", clientLocation, accessKey), data.ProxyTarget)
	addIfNonZero(m.dataBytes.WithLabelValues("p<t", "tcp", clientLocation, accessKey), data.TargetProxy)
	addIfNonZero(m.dataBytes.WithLabelValues("c<p", "tcp", clientLocation, accessKey), data.ProxyClient)
}

func (m *shadowsocksMetrics) AddTCPProbe(clientLocation, status, drainResult string, port int, data ProxyMetrics) {
	m.tcpProbes.WithLabelValues(clientLocation, strconv.Itoa(port), status, drainResult).Observe(float64(data.ClientProxy))
}

func (m *shadowsocksMetrics) AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int64, timeToCipher time.Duration) {
	m.timeToCipherMs.WithLabelValues("udp", isFound(accessKey)).Observe(timeToCipher.Seconds() * 1000)
	addIfNonZero(m.dataBytes.WithLabelValues("c>p", "udp", clientLocation, accessKey), clientProxyBytes)
	addIfNonZero(m.dataBytes.WithLabelValues("p>t", "udp", clientLocation, accessKey), proxyTargetBytes)
}

func (m *shadowsocksMetrics) AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int64) {
	addIfNonZero(m.dataBytes.WithLabelValues("p<t", "udp", clientLocation, accessKey), targetProxyBytes)
	addIfNonZero(m.dataBytes.WithLabelValues("c<p", "udp", clientLocation, accessKey), proxyClientBytes)
}

func (m *shadowsocksMetrics) AddUDPNatEntry() {
	m.udpAddedNatEntries.Inc()
}

func (m *shadowsocksMetrics) RemoveUDPNatEntry() {
	m.udpRemovedNatEntries.Inc()
}

type ProxyMetrics struct {
	ClientProxy int64
	ProxyTarget int64
	TargetProxy int64
	ProxyClient int64
}

func (m *ProxyMetrics) add(other ProxyMetrics) {
	m.ClientProxy += other.ClientProxy
	m.ProxyTarget += other.ProxyTarget
	m.TargetProxy += other.TargetProxy
	m.ProxyClient += other.ProxyClient
}

type measuredConn struct {
	onet.DuplexConn
	io.WriterTo
	readCount *int64
	io.ReaderFrom
	writeCount *int64
}

func (c *measuredConn) Read(b []byte) (int, error) {
	n, err := c.DuplexConn.Read(b)
	*c.readCount += int64(n)
	return n, err
}

func (c *measuredConn) WriteTo(w io.Writer) (int64, error) {
	n, err := io.Copy(w, c.DuplexConn)
	*c.readCount += n
	return n, err
}

func (c *measuredConn) Write(b []byte) (int, error) {
	n, err := c.DuplexConn.Write(b)
	*c.writeCount += int64(n)
	return n, err
}

func (c *measuredConn) ReadFrom(r io.Reader) (int64, error) {
	n, err := io.Copy(c.DuplexConn, r)
	*c.writeCount += n
	return n, err
}

func MeasureConn(conn onet.DuplexConn, bytesSent, bytesReceived *int64) onet.DuplexConn {
	return &measuredConn{DuplexConn: conn, writeCount: bytesSent, readCount: bytesReceived}
}

// NoOpMetrics is a fake ShadowsocksMetrics that doesn't do anything. Useful in tests
// or if you don't want to track metrics.
type NoOpMetrics struct{}

func (m *NoOpMetrics) SetBuildInfo(version string) {}
func (m *NoOpMetrics) AddTCPProbe(clientLocation, status, drainResult string, port int, data ProxyMetrics) {
}
func (m *NoOpMetrics) AddClosedTCPConnection(clientLocation, accessKey, status string, data ProxyMetrics, timeToCipher, duration time.Duration) {
}
func (m *NoOpMetrics) GetLocation(net.Addr) (string, error) {
	return "", nil
}
func (m *NoOpMetrics) SetNumAccessKeys(numKeys int, numPorts int) {}
func (m *NoOpMetrics) AddOpenTCPConnection(clientLocation string) {}
func (m *NoOpMetrics) AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration) {
}
func (m *NoOpMetrics) AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
}
func (m *NoOpMetrics) AddUDPNatEntry()    {}
func (m *NoOpMetrics) RemoveUDPNatEntry() {}
<commit_msg>Fix integer type<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package metrics

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"time"

	onet "github.com/Jigsaw-Code/outline-ss-server/net"
	geoip2 "github.com/oschwald/geoip2-golang"
	"github.com/prometheus/client_golang/prometheus"
)

// ShadowsocksMetrics registers metrics for the Shadowsocks service.
type ShadowsocksMetrics interface {
	SetBuildInfo(version string)

	GetLocation(net.Addr) (string, error)

	SetNumAccessKeys(numKeys int, numPorts int)

	// TCP metrics
	AddOpenTCPConnection(clientLocation string)
	AddClosedTCPConnection(clientLocation, accessKey, status string, data ProxyMetrics, timeToCipher, duration time.Duration)
	AddTCPProbe(clientLocation, status, drainResult string, port int, data ProxyMetrics)

	// UDP metrics
	AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration)
	AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int)
	AddUDPNatEntry()
	RemoveUDPNatEntry()
}

type shadowsocksMetrics struct {
	ipCountryDB *geoip2.Reader

	buildInfo      *prometheus.GaugeVec
	accessKeys     prometheus.Gauge
	ports          prometheus.Gauge
	dataBytes      *prometheus.CounterVec
	timeToCipherMs *prometheus.HistogramVec
	// TODO: Add time to first byte.

	tcpProbes               *prometheus.HistogramVec
	tcpOpenConnections      *prometheus.CounterVec
	tcpClosedConnections    *prometheus.CounterVec
	tcpConnectionDurationMs *prometheus.HistogramVec

	udpAddedNatEntries   prometheus.Counter
	udpRemovedNatEntries prometheus.Counter
}

func newShadowsocksMetrics(ipCountryDB *geoip2.Reader) *shadowsocksMetrics {
	return &shadowsocksMetrics{
		ipCountryDB: ipCountryDB,
		buildInfo: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Namespace: "shadowsocks",
			Name:      "build_info",
			Help:      "Information on the outline-ss-server build",
		}, []string{"version"}),
		accessKeys: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "shadowsocks",
			Name:      "keys",
			Help:      "Count of access keys",
		}),
		ports: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: "shadowsocks",
			Name:      "ports",
			Help:      "Count of open Shadowsocks ports",
		}),
		tcpOpenConnections: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: "shadowsocks",
			Subsystem: "tcp",
			Name:      "connections_opened",
			Help:      "Count of open TCP connections",
		}, []string{"location"}),
		tcpClosedConnections: prometheus.NewCounterVec(prometheus.CounterOpts{
			Namespace: "shadowsocks",
			Subsystem: "tcp",
			Name:      "connections_closed",
			Help:      "Count of closed TCP connections",
		}, []string{"location", "status", "access_key"}),
		tcpConnectionDurationMs: prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: "shadowsocks",
				Subsystem: "tcp",
				Name:      "connection_duration_ms",
				Help:      "TCP connection duration distributions.",
				Buckets: []float64{
					100,
					float64(time.Second.Milliseconds()),
					float64(time.Minute.Milliseconds()),
					float64(time.Hour.Milliseconds()),
					float64(24 * time.Hour.Milliseconds()),     // Day
					float64(7 * 24 * time.Hour.Milliseconds()), // Week
				},
			}, []string{"status"}),
		dataBytes: prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Namespace: "shadowsocks",
				Name:      "data_bytes",
				Help:      "Bytes transferred by the proxy",
			}, []string{"dir", "proto", "location", "access_key"}),
		tcpProbes: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Namespace: "shadowsocks",
			Name:      "tcp_probes",
			Buckets:   []float64{0, 48, 49, 50, 51, 52, 72, 73, 90, 91, 220, 221},
			Help:      "Histogram of number of bytes from client to proxy, for detecting possible probes",
		}, []string{"location", "port", "status", "error"}),
		timeToCipherMs: prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Namespace: "shadowsocks",
				Name:      "time_to_cipher_ms",
				Help:      "Time needed to find the cipher",
				Buckets:   []float64{0.1, 1, 10, 100, 1000},
			}, []string{"proto", "found_key"}),
		udpAddedNatEntries: prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "shadowsocks",
				Subsystem: "udp",
				Name:      "nat_entries_added",
				Help:      "Entries added to the UDP NAT table",
			}),
		udpRemovedNatEntries: prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "shadowsocks",
				Subsystem: "udp",
				Name:      "nat_entries_removed",
				Help:      "Entries removed from the UDP NAT table",
			}),
	}
}

// NewPrometheusShadowsocksMetrics constructs a metrics object that uses
// `ipCountryDB` to convert IP addresses to countries, and reports all
// metrics to Prometheus via `registerer`.  `ipCountryDB` may be nil, but
// `registerer` must not be.
func NewPrometheusShadowsocksMetrics(ipCountryDB *geoip2.Reader, registerer prometheus.Registerer) ShadowsocksMetrics {
	m := newShadowsocksMetrics(ipCountryDB)
	// TODO: Is it possible to pass where to register the collectors?
	registerer.MustRegister(m.buildInfo, m.accessKeys, m.ports, m.tcpOpenConnections, m.tcpProbes, m.tcpClosedConnections, m.tcpConnectionDurationMs,
		m.dataBytes, m.timeToCipherMs, m.udpAddedNatEntries, m.udpRemovedNatEntries)
	return m
}

const (
	errParseAddr     = "XA"
	errDbLookupError = "XD"
	localLocation    = "XL"
	unknownLocation  = "ZZ"
)

func (m *shadowsocksMetrics) SetBuildInfo(version string) {
	m.buildInfo.WithLabelValues(version).Set(1)
}

func (m *shadowsocksMetrics) GetLocation(addr net.Addr) (string, error) {
	if m.ipCountryDB == nil {
		return "", nil
	}
	hostname, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return errParseAddr, errors.New("Failed to split hostname and port")
	}
	ip := net.ParseIP(hostname)
	if ip == nil {
		return errParseAddr, errors.New("Failed to parse address as IP")
	}
	if ip.IsLoopback() {
		return localLocation, nil
	}
	if !ip.IsGlobalUnicast() {
		return localLocation, nil
	}
	record, err := m.ipCountryDB.Country(ip)
	if err != nil {
		return errDbLookupError, errors.New("IP lookup failed")
	}
	if record == nil {
		return unknownLocation, errors.New("IP lookup returned nil")
	}
	if record.Country.IsoCode == "" {
		return unknownLocation, errors.New("IP Lookup has empty ISO code")
	}
	return record.Country.IsoCode, nil
}

func (m *shadowsocksMetrics) SetNumAccessKeys(numKeys int, ports int) {
	m.accessKeys.Set(float64(numKeys))
	m.ports.Set(float64(ports))
}

func (m *shadowsocksMetrics) AddOpenTCPConnection(clientLocation string) {
	m.tcpOpenConnections.WithLabelValues(clientLocation).Inc()
}

// Converts accessKey to "true" or "false"
func isFound(accessKey string) string {
	return fmt.Sprintf("%t", accessKey != "")
}

// addIfNonZero helps avoid the creation of series that are always zero.
func addIfNonZero(counter prometheus.Counter, value float64) {
	if value > 0 {
		counter.Add(value)
	}
}

func (m *shadowsocksMetrics) AddClosedTCPConnection(clientLocation, accessKey, status string, data ProxyMetrics, timeToCipher, duration time.Duration) {
	m.tcpClosedConnections.WithLabelValues(clientLocation, status, accessKey).Inc()
	m.tcpConnectionDurationMs.WithLabelValues(status).Observe(duration.Seconds() * 1000)
	m.timeToCipherMs.WithLabelValues("tcp", isFound(accessKey)).Observe(timeToCipher.Seconds() * 1000)
	addIfNonZero(m.dataBytes.WithLabelValues("c>p", "tcp", clientLocation, accessKey), float64(data.ClientProxy))
	addIfNonZero(m.dataBytes.WithLabelValues("p>t", "tcp", clientLocation, accessKey), float64(data.ProxyTarget))
	addIfNonZero(m.dataBytes.WithLabelValues("p<t", "tcp", clientLocation, accessKey), float64(data.TargetProxy))
	addIfNonZero(m.dataBytes.WithLabelValues("c<p", "tcp", clientLocation, accessKey), float64(data.ProxyClient))
}

func (m *shadowsocksMetrics) AddTCPProbe(clientLocation, status, drainResult string, port int, data ProxyMetrics) {
	m.tcpProbes.WithLabelValues(clientLocation, strconv.Itoa(port), status, drainResult).Observe(float64(data.ClientProxy))
}

func (m *shadowsocksMetrics) AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration) {
	m.timeToCipherMs.WithLabelValues("udp", isFound(accessKey)).Observe(timeToCipher.Seconds() * 1000)
	addIfNonZero(m.dataBytes.WithLabelValues("c>p", "udp", clientLocation, accessKey), float64(clientProxyBytes))
	addIfNonZero(m.dataBytes.WithLabelValues("p>t", "udp", clientLocation, accessKey), float64(proxyTargetBytes))
}

func (m *shadowsocksMetrics) AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
	addIfNonZero(m.dataBytes.WithLabelValues("p<t", "udp", clientLocation, accessKey), float64(targetProxyBytes))
	addIfNonZero(m.dataBytes.WithLabelValues("c<p", "udp", clientLocation, accessKey), float64(proxyClientBytes))
}

func (m *shadowsocksMetrics) AddUDPNatEntry() {
	m.udpAddedNatEntries.Inc()
}

func (m *shadowsocksMetrics) RemoveUDPNatEntry() {
	m.udpRemovedNatEntries.Inc()
}

type ProxyMetrics struct {
	ClientProxy int64
	ProxyTarget int64
	TargetProxy int64
	ProxyClient int64
}

func (m *ProxyMetrics) add(other ProxyMetrics) {
	m.ClientProxy += other.ClientProxy
	m.ProxyTarget += other.ProxyTarget
	m.TargetProxy += other.TargetProxy
	m.ProxyClient += other.ProxyClient
}

type measuredConn struct {
	onet.DuplexConn
	io.WriterTo
	readCount *int64
	io.ReaderFrom
	writeCount *int64
}

func (c *measuredConn) Read(b []byte) (int, error) {
	n, err := c.DuplexConn.Read(b)
	*c.readCount += int64(n)
	return n, err
}

func (c *measuredConn) WriteTo(w io.Writer) (int64, error) {
	n, err := io.Copy(w, c.DuplexConn)
	*c.readCount += n
	return n, err
}

func (c *measuredConn) Write(b []byte) (int, error) {
	n, err := c.DuplexConn.Write(b)
	*c.writeCount += int64(n)
	return n, err
}

func (c *measuredConn) ReadFrom(r io.Reader) (int64, error) {
	n, err := io.Copy(c.DuplexConn, r)
	*c.writeCount += n
	return n, err
}

func MeasureConn(conn onet.DuplexConn, bytesSent, bytesReceived *int64) onet.DuplexConn {
	return &measuredConn{DuplexConn: conn, writeCount: bytesSent, readCount: bytesReceived}
}

// NoOpMetrics is a fake ShadowsocksMetrics that doesn't do anything. Useful in tests
// or if you don't want to track metrics.
type NoOpMetrics struct{}

func (m *NoOpMetrics) SetBuildInfo(version string) {}
func (m *NoOpMetrics) AddTCPProbe(clientLocation, status, drainResult string, port int, data ProxyMetrics) {
}
func (m *NoOpMetrics) AddClosedTCPConnection(clientLocation, accessKey, status string, data ProxyMetrics, timeToCipher, duration time.Duration) {
}
func (m *NoOpMetrics) GetLocation(net.Addr) (string, error) {
	return "", nil
}
func (m *NoOpMetrics) SetNumAccessKeys(numKeys int, numPorts int) {}
func (m *NoOpMetrics) AddOpenTCPConnection(clientLocation string) {}
func (m *NoOpMetrics) AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration) {
}
func (m *NoOpMetrics) AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
}
func (m *NoOpMetrics) AddUDPNatEntry()    {}
func (m *NoOpMetrics) RemoveUDPNatEntry() {}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/micro/cli/v2"
	goclient "github.com/micro/go-micro/v3/client"
	"github.com/micro/micro/v3/client/cli/util"
	"github.com/micro/micro/v3/cmd"
	clic "github.com/micro/micro/v3/internal/command"
	"github.com/micro/micro/v3/service/client"
	"github.com/micro/micro/v3/service/context"
	"github.com/olekukonko/tablewriter"
)

func init() {
	cmd.Register(&cli.Command{
		Name:  "network",
		Usage: "Manage the micro service network",
		Subcommands: []*cli.Command{
			{
				Name:   "connect",
				Usage:  "connect to the network. specify nodes e.g connect ip:port",
				Action: util.Print(networkConnect),
			},
			{
				Name:   "connections",
				Usage:  "List the immediate connections to the network",
				Action: util.Print(networkConnections),
			},
			{
				Name:   "graph",
				Usage:  "Get the network graph",
				Action: util.Print(networkGraph),
			},
			{
				Name:   "nodes",
				Usage:  "List nodes in the network",
				Action: util.Print(networkNodes),
			},
			{
				Name:   "routes",
				Usage:  "List network routes",
				Action: util.Print(networkRoutes),
				Flags: []cli.Flag{
					&cli.StringFlag{
						Name:  "service",
						Usage: "Filter by service",
					},
					&cli.StringFlag{
						Name:  "address",
						Usage: "Filter by address",
					},
					&cli.StringFlag{
						Name:  "gateway",
						Usage: "Filter by gateway",
					},
					&cli.StringFlag{
						Name:  "router",
						Usage: "Filter by router",
					},
					&cli.StringFlag{
						Name:  "network",
						Usage: "Filter by network",
					},
				},
			},
			{
				Name:   "services",
				Usage:  "Get the network services",
				Action: util.Print(networkServices),
			},
			// TODO: duplicates call. Move so we reuse same stuff.
			{
				Name:   "call",
				Usage:  "Call a service e.g micro call greeter Say.Hello '{\"name\": \"John\"}",
				Action: util.Print(netCall),
				Flags: []cli.Flag{
					&cli.StringFlag{
						Name:    "address",
						Usage:   "Set the address of the service instance to call",
						EnvVars: []string{"MICRO_ADDRESS"},
					},
					&cli.StringFlag{
						Name:    "output, o",
						Usage:   "Set the output format; json (default), raw",
						EnvVars: []string{"MICRO_OUTPUT"},
					},
					&cli.StringSliceFlag{
						Name:    "metadata",
						Usage:   "A list of key-value pairs to be forwarded as metadata",
						EnvVars: []string{"MICRO_METADATA"},
					},
				},
			},
		},
	})
}

func networkConnect(c *cli.Context, args []string) ([]byte, error) {
	if len(args) == 0 {
		return nil, nil
	}

	request := map[string]interface{}{
		"nodes": []interface{}{
			map[string]interface{}{
				"address": args[0],
			},
		},
	}

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Connect", request, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp)
	if err != nil {
		return nil, err
	}

	b, _ := json.MarshalIndent(rsp, "", "\t")
	return b, nil
}

func networkConnections(c *cli.Context, args []string) ([]byte, error) {

	request := map[string]interface{}{
		"depth": 1,
	}

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Graph", request, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp)
	if err != nil {
		return nil, err
	}

	if rsp["root"] == nil {
		return nil, nil
	}

	peers := rsp["root"].(map[string]interface{})["peers"]

	if peers == nil {
		return nil, nil
	}

	b := bytes.NewBuffer(nil)
	table := tablewriter.NewWriter(b)
	table.SetHeader([]string{"NODE", "ADDRESS"})

	// root node
	for _, n := range peers.([]interface{}) {
		node := n.(map[string]interface{})["node"].(map[string]interface{})
		strEntry := []string{
			fmt.Sprintf("%s", node["id"]),
			fmt.Sprintf("%s", node["address"]),
		}
		table.Append(strEntry)
	}

	// render table into b
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()

	return b.Bytes(), nil
}

func networkGraph(c *cli.Context, args []string) ([]byte, error) {

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Graph", map[string]interface{}{}, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp)
	if err != nil {
		return nil, err
	}

	b, _ := json.MarshalIndent(rsp, "", "\t")
	return b, nil
}

func networkNodes(c *cli.Context, args []string) ([]byte, error) {

	var rsp map[string]interface{}

	// TODO: change to list nodes
	req := client.NewRequest("network", "Network.Nodes", map[string]interface{}{}, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp)
	if err != nil {
		return nil, err
	}

	// return if nil
	if rsp["nodes"] == nil {
		return nil, nil
	}

	b := bytes.NewBuffer(nil)
	table := tablewriter.NewWriter(b)
	table.SetHeader([]string{"ID", "ADDRESS"})

	// get nodes

	if rsp["nodes"] != nil {
		// root node
		for _, n := range rsp["nodes"].([]interface{}) {
			node := n.(map[string]interface{})
			strEntry := []string{
				fmt.Sprintf("%s", node["id"]),
				fmt.Sprintf("%s", node["address"]),
			}
			table.Append(strEntry)
		}
	}

	// render table into b
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()

	return b.Bytes(), nil
}

func networkRoutes(c *cli.Context, args []string) ([]byte, error) {

	query := map[string]string{}

	for _, filter := range []string{"service", "address", "gateway", "router", "network"} {
		if v := c.String(filter); len(v) > 0 {
			query[filter] = v
		}
	}

	request := map[string]interface{}{
		"query": query,
	}

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Routes", request, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	if len(rsp) == 0 {
		return []byte(``), nil
	}

	b := bytes.NewBuffer(nil)
	table := tablewriter.NewWriter(b)
	table.SetHeader([]string{"SERVICE", "ADDRESS", "GATEWAY", "ROUTER", "NETWORK", "METRIC", "LINK"})

	routes := rsp["routes"].([]interface{})

	val := func(v interface{}) string {
		if v == nil {
			return ""
		}
		return v.(string)
	}

	var sortedRoutes [][]string

	for _, r := range routes {
		route := r.(map[string]interface{})
		service := route["service"]
		address := route["address"]
		gateway := val(route["gateway"])
		router := route["router"]
		network := route["network"]
		link := route["link"]
		metric := route["metric"]

		var metInt int64
		if metric != nil {
			metInt, _ = strconv.ParseInt(route["metric"].(string), 10, 64)
		}

		// set max int64 metric to infinity
		if metInt == math.MaxInt64 {
			metric = "∞"
		} else {
			metric = fmt.Sprintf("%d", metInt)
		}

		sortedRoutes = append(sortedRoutes, []string{
			fmt.Sprintf("%s", service),
			fmt.Sprintf("%s", address),
			fmt.Sprintf("%s", gateway),
			fmt.Sprintf("%s", router),
			fmt.Sprintf("%s", network),
			fmt.Sprintf("%s", metric),
			fmt.Sprintf("%s", link),
		})
	}

	sort.Slice(sortedRoutes, func(i, j int) bool { return sortedRoutes[i][0] < sortedRoutes[j][0] })

	table.AppendBulk(sortedRoutes)
	// render table into b
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()

	return b.Bytes(), nil
}

func networkServices(c *cli.Context, args []string) ([]byte, error) {

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Services", map[string]interface{}{}, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	if len(rsp) == 0 || rsp["services"] == nil {
		return []byte(``), nil
	}

	rspSrv := rsp["services"].([]interface{})

	var services []string

	for _, service := range rspSrv {
		services = append(services, service.(string))
	}

	sort.Strings(services)

	return []byte(strings.Join(services, "\n")), nil
}

// netCall calls services through the network
func netCall(c *cli.Context, args []string) ([]byte, error) {
	os.Setenv("MICRO_PROXY", "network")
	return clic.CallService(c, args)
}
<commit_msg>Fixing auth of network commands (#1301)<commit_after>package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"sort"
	"strconv"
	"strings"

	"github.com/micro/cli/v2"
	goclient "github.com/micro/go-micro/v3/client"
	"github.com/micro/micro/v3/client/cli/util"
	"github.com/micro/micro/v3/cmd"
	clic "github.com/micro/micro/v3/internal/command"
	"github.com/micro/micro/v3/service/client"
	"github.com/micro/micro/v3/service/context"
	"github.com/olekukonko/tablewriter"
)

func init() {
	cmd.Register(&cli.Command{
		Name:  "network",
		Usage: "Manage the micro service network",
		Subcommands: []*cli.Command{
			{
				Name:   "connect",
				Usage:  "connect to the network. specify nodes e.g connect ip:port",
				Action: util.Print(networkConnect),
			},
			{
				Name:   "connections",
				Usage:  "List the immediate connections to the network",
				Action: util.Print(networkConnections),
			},
			{
				Name:   "graph",
				Usage:  "Get the network graph",
				Action: util.Print(networkGraph),
			},
			{
				Name:   "nodes",
				Usage:  "List nodes in the network",
				Action: util.Print(networkNodes),
			},
			{
				Name:   "routes",
				Usage:  "List network routes",
				Action: util.Print(networkRoutes),
				Flags: []cli.Flag{
					&cli.StringFlag{
						Name:  "service",
						Usage: "Filter by service",
					},
					&cli.StringFlag{
						Name:  "address",
						Usage: "Filter by address",
					},
					&cli.StringFlag{
						Name:  "gateway",
						Usage: "Filter by gateway",
					},
					&cli.StringFlag{
						Name:  "router",
						Usage: "Filter by router",
					},
					&cli.StringFlag{
						Name:  "network",
						Usage: "Filter by network",
					},
				},
			},
			{
				Name:   "services",
				Usage:  "Get the network services",
				Action: util.Print(networkServices),
			},
			// TODO: duplicates call. Move so we reuse same stuff.
			{
				Name:   "call",
				Usage:  "Call a service e.g micro call greeter Say.Hello '{\"name\": \"John\"}",
				Action: util.Print(netCall),
				Flags: []cli.Flag{
					&cli.StringFlag{
						Name:    "address",
						Usage:   "Set the address of the service instance to call",
						EnvVars: []string{"MICRO_ADDRESS"},
					},
					&cli.StringFlag{
						Name:    "output, o",
						Usage:   "Set the output format; json (default), raw",
						EnvVars: []string{"MICRO_OUTPUT"},
					},
					&cli.StringSliceFlag{
						Name:    "metadata",
						Usage:   "A list of key-value pairs to be forwarded as metadata",
						EnvVars: []string{"MICRO_METADATA"},
					},
				},
			},
		},
	})
}

func networkConnect(c *cli.Context, args []string) ([]byte, error) {
	if len(args) == 0 {
		return nil, nil
	}

	request := map[string]interface{}{
		"nodes": []interface{}{
			map[string]interface{}{
				"address": args[0],
			},
		},
	}

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Connect", request, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	b, _ := json.MarshalIndent(rsp, "", "\t")
	return b, nil
}

func networkConnections(c *cli.Context, args []string) ([]byte, error) {

	request := map[string]interface{}{
		"depth": 1,
	}

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Graph", request, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	if rsp["root"] == nil {
		return nil, nil
	}

	peers := rsp["root"].(map[string]interface{})["peers"]

	if peers == nil {
		return nil, nil
	}

	b := bytes.NewBuffer(nil)
	table := tablewriter.NewWriter(b)
	table.SetHeader([]string{"NODE", "ADDRESS"})

	// root node
	for _, n := range peers.([]interface{}) {
		node := n.(map[string]interface{})["node"].(map[string]interface{})
		strEntry := []string{
			fmt.Sprintf("%s", node["id"]),
			fmt.Sprintf("%s", node["address"]),
		}
		table.Append(strEntry)
	}

	// render table into b
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()

	return b.Bytes(), nil
}

func networkGraph(c *cli.Context, args []string) ([]byte, error) {

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Graph", map[string]interface{}{}, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	b, _ := json.MarshalIndent(rsp, "", "\t")
	return b, nil
}

func networkNodes(c *cli.Context, args []string) ([]byte, error) {

	var rsp map[string]interface{}

	// TODO: change to list nodes
	req := client.NewRequest("network", "Network.Nodes", map[string]interface{}{}, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	// return if nil
	if rsp["nodes"] == nil {
		return nil, nil
	}

	b := bytes.NewBuffer(nil)
	table := tablewriter.NewWriter(b)
	table.SetHeader([]string{"ID", "ADDRESS"})

	// get nodes

	if rsp["nodes"] != nil {
		// root node
		for _, n := range rsp["nodes"].([]interface{}) {
			node := n.(map[string]interface{})
			strEntry := []string{
				fmt.Sprintf("%s", node["id"]),
				fmt.Sprintf("%s", node["address"]),
			}
			table.Append(strEntry)
		}
	}

	// render table into b
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()

	return b.Bytes(), nil
}

func networkRoutes(c *cli.Context, args []string) ([]byte, error) {

	query := map[string]string{}

	for _, filter := range []string{"service", "address", "gateway", "router", "network"} {
		if v := c.String(filter); len(v) > 0 {
			query[filter] = v
		}
	}

	request := map[string]interface{}{
		"query": query,
	}

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Routes", request, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	if len(rsp) == 0 {
		return []byte(``), nil
	}

	b := bytes.NewBuffer(nil)
	table := tablewriter.NewWriter(b)
	table.SetHeader([]string{"SERVICE", "ADDRESS", "GATEWAY", "ROUTER", "NETWORK", "METRIC", "LINK"})

	routes := rsp["routes"].([]interface{})

	val := func(v interface{}) string {
		if v == nil {
			return ""
		}
		return v.(string)
	}

	var sortedRoutes [][]string

	for _, r := range routes {
		route := r.(map[string]interface{})
		service := route["service"]
		address := route["address"]
		gateway := val(route["gateway"])
		router := route["router"]
		network := route["network"]
		link := route["link"]
		metric := route["metric"]

		var metInt int64
		if metric != nil {
			metInt, _ = strconv.ParseInt(route["metric"].(string), 10, 64)
		}

		// set max int64 metric to infinity
		if metInt == math.MaxInt64 {
			metric = "∞"
		} else {
			metric = fmt.Sprintf("%d", metInt)
		}

		sortedRoutes = append(sortedRoutes, []string{
			fmt.Sprintf("%s", service),
			fmt.Sprintf("%s", address),
			fmt.Sprintf("%s", gateway),
			fmt.Sprintf("%s", router),
			fmt.Sprintf("%s", network),
			fmt.Sprintf("%s", metric),
			fmt.Sprintf("%s", link),
		})
	}

	sort.Slice(sortedRoutes, func(i, j int) bool { return sortedRoutes[i][0] < sortedRoutes[j][0] })

	table.AppendBulk(sortedRoutes)
	// render table into b
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()

	return b.Bytes(), nil
}

func networkServices(c *cli.Context, args []string) ([]byte, error) {

	var rsp map[string]interface{}

	req := client.NewRequest("network", "Network.Services", map[string]interface{}{}, goclient.WithContentType("application/json"))
	err := client.Call(context.DefaultContext, req, &rsp, goclient.WithAuthToken())
	if err != nil {
		return nil, err
	}

	if len(rsp) == 0 || rsp["services"] == nil {
		return []byte(``), nil
	}

	rspSrv := rsp["services"].([]interface{})

	var services []string

	for _, service := range rspSrv {
		services = append(services, service.(string))
	}

	sort.Strings(services)

	return []byte(strings.Join(services, "\n")), nil
}

// netCall calls services through the network
func netCall(c *cli.Context, args []string) ([]byte, error) {
	os.Setenv("MICRO_PROXY", "network")
	return clic.CallService(c, args)
}
<|endoftext|>


<commit_before>package ssdp

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"time"
)

// Service is discovered service.
type Service struct {
	// Type is a property of "ST"
	Type string

	// USN is a property of "USN"
	USN string

	// Location is a property of "LOCATION"
	Location string

	// Server is a property of "SERVER"
	Server string

	rawHeader http.Header
	maxAge    *int
}

var rxMaxAge = regexp.MustCompile(`\bmax-age\s*=\s*(\d+)\b`)

func extractMaxAge(s string, value int) int {
	v := value
	if m := rxMaxAge.FindStringSubmatch(s); m != nil {
		i64, err := strconv.ParseInt(m[1], 10, 32)
		if err == nil {
			v = int(i64)
		}
	}
	return v
}

// MaxAge extracts "max-age" value from "CACHE-CONTROL" property.
func (s *Service) MaxAge() int {
	if s.maxAge == nil {
		s.maxAge = new(int)
		*s.maxAge = extractMaxAge(s.rawHeader.Get("CACHE-CONTROL"), -1)
	}
	return *s.maxAge
}

// Header returns all properties in response of search.
func (s *Service) Header() http.Header {
	return s.rawHeader
}

const (
	// All is a search type to search all services and devices.
	All = "ssdp:all"

	// RootDevice is a search type to search UPnP root devices.
	RootDevice = "upnp:rootdevice"
)

// Search searchs services by SSDP.
func Search(searchType string, waitSec int, localAddr string) ([]Service, error) {
	// dial multicast UDP packet.
	conn, err := multicastListen(localAddr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	logf("search on %s", conn.LocalAddr().String())

	// send request.
	msg, err := buildSearch(ssdpAddrIPv4, searchType, waitSec)
	if err != nil {
		return nil, err
	}
	if _, err := conn.WriteTo(msg, ssdpAddrIPv4); err != nil {
		return nil, err
	}

	// wait response.
	var list []Service
	h := func(a net.Addr, d []byte) error {
		srv, err := parseService(a, d)
		if err != nil {
			logf("invalid search response from %s: %s", a.String(), err)
			return nil
		}
		list = append(list, *srv)
		logf("search response from %s: %s", a.String(), srv.USN)
		return nil
	}
	d := time.Second * time.Duration(waitSec)
	if err := conn.readPackets(d, h); err != nil {
		return nil, err
	}

	return list, err
}

func buildSearch(raddr net.Addr, searchType string, waitSec int) ([]byte, error) {
	b := new(bytes.Buffer)
	// FIXME: error should be checked.
	b.WriteString("M-SEARCH * HTTP/1.1\r\n")
	fmt.Fprintf(b, "HOST: %s\r\n", raddr.String())
	fmt.Fprintf(b, "MAN: %q\r\n", "ssdp:discover")
	fmt.Fprintf(b, "MX: %d\r\n", waitSec)
	fmt.Fprintf(b, "ST: %s\r\n", searchType)
	b.WriteString("\r\n")
	return b.Bytes(), nil
}

var (
	errWithoutHTTPPrefix = errors.New("without HTTP prefix")
)

// Actually is a double new line
var newLine = []byte{'\r', '\n', '\r', '\n'}

func parseService(addr net.Addr, data []byte) (*Service, error) {
	if !bytes.HasPrefix(data, []byte("HTTP")) {
		return nil, errWithoutHTTPPrefix
	}
	// Add newline to workaround buggy SSDP responses
	if !bytes.HasSuffix(data, newLine) {
		// why we should't use append() for this purpose:
		// https://play.golang.org/p/IM1pONW9lqm
		data = bytes.Join([][]byte{data, newLine}, nil)
	}
	resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(data)), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return &Service{
		Type:      resp.Header.Get("ST"),
		USN:       resp.Header.Get("USN"),
		Location:  resp.Header.Get("LOCATION"),
		Server:    resp.Header.Get("SERVER"),
		rawHeader: resp.Header,
	}, nil
}
<commit_msg>arrange wordings<commit_after>package ssdp

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"time"
)

// Service is discovered service.
type Service struct {
	// Type is a property of "ST"
	Type string

	// USN is a property of "USN"
	USN string

	// Location is a property of "LOCATION"
	Location string

	// Server is a property of "SERVER"
	Server string

	rawHeader http.Header
	maxAge    *int
}

var rxMaxAge = regexp.MustCompile(`\bmax-age\s*=\s*(\d+)\b`)

func extractMaxAge(s string, value int) int {
	v := value
	if m := rxMaxAge.FindStringSubmatch(s); m != nil {
		i64, err := strconv.ParseInt(m[1], 10, 32)
		if err == nil {
			v = int(i64)
		}
	}
	return v
}

// MaxAge extracts "max-age" value from "CACHE-CONTROL" property.
func (s *Service) MaxAge() int {
	if s.maxAge == nil {
		s.maxAge = new(int)
		*s.maxAge = extractMaxAge(s.rawHeader.Get("CACHE-CONTROL"), -1)
	}
	return *s.maxAge
}

// Header returns all properties in response of search.
func (s *Service) Header() http.Header {
	return s.rawHeader
}

const (
	// All is a search type to search all services and devices.
	All = "ssdp:all"

	// RootDevice is a search type to search UPnP root devices.
	RootDevice = "upnp:rootdevice"
)

// Search searchs services by SSDP.
func Search(searchType string, waitSec int, localAddr string) ([]Service, error) {
	// dial multicast UDP packet.
	conn, err := multicastListen(localAddr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	logf("search on %s", conn.LocalAddr().String())

	// send request.
	msg, err := buildSearch(ssdpAddrIPv4, searchType, waitSec)
	if err != nil {
		return nil, err
	}
	if _, err := conn.WriteTo(msg, ssdpAddrIPv4); err != nil {
		return nil, err
	}

	// wait response.
	var list []Service
	h := func(a net.Addr, d []byte) error {
		srv, err := parseService(a, d)
		if err != nil {
			logf("invalid search response from %s: %s", a.String(), err)
			return nil
		}
		list = append(list, *srv)
		logf("search response from %s: %s", a.String(), srv.USN)
		return nil
	}
	d := time.Second * time.Duration(waitSec)
	if err := conn.readPackets(d, h); err != nil {
		return nil, err
	}

	return list, err
}

func buildSearch(raddr net.Addr, searchType string, waitSec int) ([]byte, error) {
	b := new(bytes.Buffer)
	// FIXME: error should be checked.
	b.WriteString("M-SEARCH * HTTP/1.1\r\n")
	fmt.Fprintf(b, "HOST: %s\r\n", raddr.String())
	fmt.Fprintf(b, "MAN: %q\r\n", "ssdp:discover")
	fmt.Fprintf(b, "MX: %d\r\n", waitSec)
	fmt.Fprintf(b, "ST: %s\r\n", searchType)
	b.WriteString("\r\n")
	return b.Bytes(), nil
}

var (
	errWithoutHTTPPrefix = errors.New("without HTTP prefix")
)

var endOfHeader = []byte{'\r', '\n', '\r', '\n'}

func parseService(addr net.Addr, data []byte) (*Service, error) {
	if !bytes.HasPrefix(data, []byte("HTTP")) {
		return nil, errWithoutHTTPPrefix
	}
	// Complement newlines on tail of header for buggy SSDP responses.
	if !bytes.HasSuffix(data, endOfHeader) {
		// why we should't use append() for this purpose:
		// https://play.golang.org/p/IM1pONW9lqm
		data = bytes.Join([][]byte{data, endOfHeader}, nil)
	}
	resp, err := http.ReadResponse(bufio.NewReader(bytes.NewReader(data)), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return &Service{
		Type:      resp.Header.Get("ST"),
		USN:       resp.Header.Get("USN"),
		Location:  resp.Header.Get("LOCATION"),
		Server:    resp.Header.Get("SERVER"),
		rawHeader: resp.Header,
	}, nil
}
<|endoftext|>


<commit_before>package chevalier

import (
	"encoding/json"
	"github.com/mattbaird/elastigo/api"
	es "github.com/mattbaird/elastigo/core"
	"github.com/mattbaird/elastigo/search"
	"log"
	"fmt"
	"strings"
	"time"
	"errors"
)

// QueryEngine presents an interface for running queries for sources
// against Elasticsearch.
type QueryEngine struct {
	indexName      string
	dataType       string
	nSources       int64
	updateInterval time.Duration
}

// NewQueryEngine initializes a QueryEngine with the supplied
// Elasticsearch metadata. indexName and dataType can be anything as
// long as they're consistent.
func NewQueryEngine(host, indexName, dataType string) *QueryEngine {
	e := new(QueryEngine)
	e.indexName = indexName
	e.dataType = dataType
	api.Domain = host
	e.updateSourceCount()
	e.updateInterval = time.Second * 10
	go e.updateForever()
	return e
}

// sanitizeField takes an input field specifier (from a
// SourceRequest_Tag) and munges it so Elasticsearch likes it more. In
// particular, it makes single-wildcard queries work correctly.
//
// FIXME: mid-field wildcards still aren't working.
func (e *QueryEngine) sanitizeField(f string) string {
	f = strings.TrimSpace(f)
	if f == "*" {
		return fmt.Sprintf("%s._all", e.dataType)
	}
	f = fmt.Sprintf("%s.%s", e.dataType, f)
	return f
}

func (e *QueryEngine) sanitizeTag(field, value string) (string, string) {
	// * is normally in this list, but is not included here because
	// we want it to act as a wildcard.
	// Also, this can be made a lot faster.
	reservedChars := `+ - && || ! ( ) { } [ ] ^ " ~ ? : \ /`
	for _, char := range strings.Split(reservedChars, " ") {
		escapedChar := fmt.Sprintf(field, "\\%s", char)
		field = strings.Replace(field, char, escapedChar, -1)
		value = strings.Replace(value, char, escapedChar, -1)
	}
	field = e.sanitizeField(field)
	return field, value
}

// buildTagQuery constructs an elastigo query object (search.QueryDsl)
// from a SourceRequest_Tag, designed to be plugged into a
// query-string-type[0] query later on. Returns error on empty query.
//
// [0]: http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html
func (e *QueryEngine) buildTagQuery(tag *SourceRequest_Tag) (*search.QueryDsl, error) {
	field, value := e.sanitizeTag(*tag.Field, *tag.Value)
	// Don't bother running empty queries.
	if value == "" {
		return nil, errors.New("empty query string")
	}
	qs := new(search.QueryString)
	qs.Fields = make([]string, 0)
	qs.Fields = append(qs.Fields, field)
	qs.Query = value
	q := search.Query().Qs(qs)
	return q, nil
}

// SourceQuery is a multi-level map type representing an Elasticsearch
// query-string-type query. Suitable for marshalling as JSON and feeding
// to Elasticsearch.
type SourceQuery map[string]interface{}

func (e *QueryEngine) getStartResult(req *SourceRequest) int64 {
	startPage := req.GetStartPage()
	pageSize := req.GetSourcesPerPage()
	if pageSize == 0 {
		return int64(0)
	}
	return startPage * pageSize
}

func (e *QueryEngine) getResultCount(req *SourceRequest) int64 {
	pageSize := req.GetSourcesPerPage()
	if pageSize <= 0 {
		return e.nSources
	}
	return pageSize
}

// BuildQuery takes a SourceRequest and turns it into a multi-level
// map suitable for marshalling to JSON and sending to Elasticsearch.
func (e *QueryEngine) BuildQuery(req *SourceRequest) (SourceQuery,error) {
	_ = search.Search(e.indexName).Type(e.dataType)
	tags := req.GetTags()
	tagQueries := make([]*search.QueryDsl, 0)
	for _, tag := range tags {
		q, err := e.buildTagQuery(tag)
		if q != nil && err == nil {
			tagQueries = append(tagQueries, q)
		}
	}
	if len(tagQueries) == 0 {
		return nil, errors.New("No valid query strings found.")
	}
	fromResult := e.getStartResult(req)
	resultCount := e.getResultCount(req)
	query := map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": tagQueries,
			},
		},
		"from": fromResult,
		"size": resultCount,
	}
	return SourceQuery(query), nil
}

// runSourceRequest takes a request object and returns an elastigo-type
// (i.e., intermediate) result.
func (e *QueryEngine) runSourceRequest(req *SourceRequest) (*es.SearchResult, error) {
	q, err := e.BuildQuery(req)
	if err != nil {
		return nil, err
	}
	res, err := es.SearchRequest(false, e.indexName, e.dataType, q, "", 0)
	return &res, err
}

// GetSources takes a request object and returns the DataSourceBurst of
// the sources it gets back from Elasticsearch. If error is not nil,
// then a valid DataSourceBurst will still be returned, with the Error
// field set.
func (e *QueryEngine) GetSources(req *SourceRequest) (*DataSourceBurst, error) {
	res, err := e.runSourceRequest(req)
	if err != nil {
		msg := fmt.Sprintf("Request error: %v", err)
		burst := new(DataSourceBurst)
		burst.Error = &msg
		return burst, err
	}
	sources := make([]*DataSource, len(res.Hits.Hits))
	for i, hit := range res.Hits.Hits {
		source := new(ElasticsearchSource)
		err = json.Unmarshal(hit.Source, source)
		if err != nil {
			msg := fmt.Sprintf("Response decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, err
		}
		sources[i] = source.Unmarshal()
	}
	burst := BuildSourceBurst(sources)
	return burst, nil
}

// FmtResult returns a string from a SearchResult by interpreting it in
// the most naive manner possible. For debugging.
func FmtResult(result *es.SearchResult) []string {
	results := make([]string, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		results[i] = string(hit.Source[:])
	}
	return results
}

// updateSourceCount updates our running total of documents-in-index
// (by asking Elasticsearch).
func (e *QueryEngine) updateSourceCount() error {
	resp, err := es.Count(false, e.indexName, e.dataType)
	e.nSources = int64(resp.Count)
	return err
}

// updateForever updates the source counter on a regular basis.
func (e *QueryEngine) updateForever() {
	for true {
		time.Sleep(e.updateInterval)
		err := e.updateSourceCount()
		if err != nil {
			log.Printf("Error updating source count: %v", err)
		}
	}
}
<commit_msg>Unfuck Sprintf invocation<commit_after>package chevalier

import (
	"encoding/json"
	"github.com/mattbaird/elastigo/api"
	es "github.com/mattbaird/elastigo/core"
	"github.com/mattbaird/elastigo/search"
	"log"
	"fmt"
	"strings"
	"time"
	"errors"
)

// QueryEngine presents an interface for running queries for sources
// against Elasticsearch.
type QueryEngine struct {
	indexName      string
	dataType       string
	nSources       int64
	updateInterval time.Duration
}

// NewQueryEngine initializes a QueryEngine with the supplied
// Elasticsearch metadata. indexName and dataType can be anything as
// long as they're consistent.
func NewQueryEngine(host, indexName, dataType string) *QueryEngine {
	e := new(QueryEngine)
	e.indexName = indexName
	e.dataType = dataType
	api.Domain = host
	e.updateSourceCount()
	e.updateInterval = time.Second * 10
	go e.updateForever()
	return e
}

// sanitizeField takes an input field specifier (from a
// SourceRequest_Tag) and munges it so Elasticsearch likes it more. In
// particular, it makes single-wildcard queries work correctly.
//
// FIXME: mid-field wildcards still aren't working.
func (e *QueryEngine) sanitizeField(f string) string {
	f = strings.TrimSpace(f)
	if f == "*" {
		return fmt.Sprintf("%s._all", e.dataType)
	}
	f = fmt.Sprintf("%s.%s", e.dataType, f)
	return f
}

func (e *QueryEngine) sanitizeTag(field, value string) (string, string) {
	// * is normally in this list, but is not included here because
	// we want it to act as a wildcard.
	// Also, this can be made a lot faster.
	reservedChars := `+ - && || ! ( ) { } [ ] ^ " ~ ? : \ /`
	for _, char := range strings.Split(reservedChars, " ") {
		escapedChar := fmt.Sprintf("\\%s", char)
		field = strings.Replace(field, char, escapedChar, -1)
		value = strings.Replace(value, char, escapedChar, -1)
	}
	field = e.sanitizeField(field)
	return field, value
}

// buildTagQuery constructs an elastigo query object (search.QueryDsl)
// from a SourceRequest_Tag, designed to be plugged into a
// query-string-type[0] query later on. Returns error on empty query.
//
// [0]: http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html
func (e *QueryEngine) buildTagQuery(tag *SourceRequest_Tag) (*search.QueryDsl, error) {
	field, value := e.sanitizeTag(*tag.Field, *tag.Value)
	// Don't bother running empty queries.
	if value == "" {
		return nil, errors.New("empty query string")
	}
	qs := new(search.QueryString)
	qs.Fields = make([]string, 0)
	qs.Fields = append(qs.Fields, field)
	qs.Query = value
	q := search.Query().Qs(qs)
	return q, nil
}

// SourceQuery is a multi-level map type representing an Elasticsearch
// query-string-type query. Suitable for marshalling as JSON and feeding
// to Elasticsearch.
type SourceQuery map[string]interface{}

func (e *QueryEngine) getStartResult(req *SourceRequest) int64 {
	startPage := req.GetStartPage()
	pageSize := req.GetSourcesPerPage()
	if pageSize == 0 {
		return int64(0)
	}
	return startPage * pageSize
}

func (e *QueryEngine) getResultCount(req *SourceRequest) int64 {
	pageSize := req.GetSourcesPerPage()
	if pageSize <= 0 {
		return e.nSources
	}
	return pageSize
}

// BuildQuery takes a SourceRequest and turns it into a multi-level
// map suitable for marshalling to JSON and sending to Elasticsearch.
func (e *QueryEngine) BuildQuery(req *SourceRequest) (SourceQuery,error) {
	_ = search.Search(e.indexName).Type(e.dataType)
	tags := req.GetTags()
	tagQueries := make([]*search.QueryDsl, 0)
	for _, tag := range tags {
		q, err := e.buildTagQuery(tag)
		if q != nil && err == nil {
			tagQueries = append(tagQueries, q)
		}
	}
	if len(tagQueries) == 0 {
		return nil, errors.New("No valid query strings found.")
	}
	fromResult := e.getStartResult(req)
	resultCount := e.getResultCount(req)
	query := map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": tagQueries,
			},
		},
		"from": fromResult,
		"size": resultCount,
	}
	return SourceQuery(query), nil
}

// runSourceRequest takes a request object and returns an elastigo-type
// (i.e., intermediate) result.
func (e *QueryEngine) runSourceRequest(req *SourceRequest) (*es.SearchResult, error) {
	q, err := e.BuildQuery(req)
	if err != nil {
		return nil, err
	}
	res, err := es.SearchRequest(false, e.indexName, e.dataType, q, "", 0)
	return &res, err
}

// GetSources takes a request object and returns the DataSourceBurst of
// the sources it gets back from Elasticsearch. If error is not nil,
// then a valid DataSourceBurst will still be returned, with the Error
// field set.
func (e *QueryEngine) GetSources(req *SourceRequest) (*DataSourceBurst, error) {
	res, err := e.runSourceRequest(req)
	if err != nil {
		msg := fmt.Sprintf("Request error: %v", err)
		burst := new(DataSourceBurst)
		burst.Error = &msg
		return burst, err
	}
	sources := make([]*DataSource, len(res.Hits.Hits))
	for i, hit := range res.Hits.Hits {
		source := new(ElasticsearchSource)
		err = json.Unmarshal(hit.Source, source)
		if err != nil {
			msg := fmt.Sprintf("Response decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, err
		}
		sources[i] = source.Unmarshal()
	}
	burst := BuildSourceBurst(sources)
	return burst, nil
}

// FmtResult returns a string from a SearchResult by interpreting it in
// the most naive manner possible. For debugging.
func FmtResult(result *es.SearchResult) []string {
	results := make([]string, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		results[i] = string(hit.Source[:])
	}
	return results
}

// updateSourceCount updates our running total of documents-in-index
// (by asking Elasticsearch).
func (e *QueryEngine) updateSourceCount() error {
	resp, err := es.Count(false, e.indexName, e.dataType)
	e.nSources = int64(resp.Count)
	return err
}

// updateForever updates the source counter on a regular basis.
func (e *QueryEngine) updateForever() {
	for true {
		time.Sleep(e.updateInterval)
		err := e.updateSourceCount()
		if err != nil {
			log.Printf("Error updating source count: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package chevalier

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/mattbaird/elastigo/api"
	es "github.com/mattbaird/elastigo/core"
	"github.com/mattbaird/elastigo/search"
	"log"
	"strings"
	"time"
)

// QueryEngine presents an interface for running queries for sources
// against Elasticsearch.
type QueryEngine struct {
	indexName      string
	metaIndex      string
	originType     string
	dataType       string
	nSources       int64
	updateInterval time.Duration
}

// NewQueryEngine initializes a QueryEngine with the supplied
// Elasticsearch metadata. indexName and dataType can be anything as
// long as they're consistent.
func NewQueryEngine(host, indexName, dataType, metaIndex string) *QueryEngine {
	e := new(QueryEngine)
	e.indexName = indexName
	e.metaIndex = metaIndex
	e.dataType = dataType
	e.originType = "chevalier_origin"
	api.Domain = host
	e.updateSourceCount()
	e.updateInterval = time.Second * 10
	go e.updateForever()
	return e
}

// sanitizeField takes an input field specifier (from a
// SourceRequest_Tag) and munges it so Elasticsearch likes it more. In
// particular, it makes single-wildcard queries work correctly.
//
// FIXME: mid-field wildcards still aren't working.
func (e *QueryEngine) sanitizeField(f string) string {
	f = strings.TrimSpace(f)
	if f == "*" {
		return fmt.Sprintf("%s._all", e.dataType)
	}
	f = fmt.Sprintf("%s.%s", e.dataType, f)
	return f
}

// sanitizeTag escapes assorted Elasticsearch wildcards.
//
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters
func (e *QueryEngine) sanitizeTag(field, value string) (string, string) {
	// * is normally in this list, but is not included here because
	// we want it to act as a wildcard.
	// Also, this can be made a lot faster.
	reservedChars := `\ + - && || ! ( ) { } [ ] ^ " ~ ? : /`
	for _, char := range strings.Split(reservedChars, " ") {
		escapedChar := fmt.Sprintf(`\%s`, char)
		field = strings.Replace(field, char, escapedChar, -1)
		value = strings.Replace(value, char, escapedChar, -1)
	}
	field = e.sanitizeField(field)
	return field, value
}

// buildTagQuery constructs an elastigo query object (search.QueryDsl)
// from a SourceRequest_Tag, designed to be plugged into a
// query-string-type[0] query later on. Returns error on empty query.
//
// [0]: http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html
func (e *QueryEngine) buildTagQuery(tag *SourceRequest_Tag) (*search.QueryDsl, error) {
	field, value := e.sanitizeTag(*tag.Field, *tag.Value)
	// Don't bother running empty queries.
	if value == "" {
		return nil, errors.New("empty query string")
	}
	qs := new(search.QueryString)
	qs.Fields = make([]string, 0)
	qs.Fields = append(qs.Fields, field)
	qs.Query = value
	q := search.Query().Qs(qs)
	return q, nil
}

// buildOriginQuery returns the origin-matching part of the complete
// query.
func (e *QueryEngine) buildOriginQuery(origin string) *search.QueryDsl {
	qs := new(search.QueryString)
	qs.Fields = make([]string, 1)
	qs.Fields[0] = "Origin"
	qs.Query = origin
	q := search.Query().Qs(qs)
	return q
}

// SourceQuery is a multi-level map type representing an Elasticsearch
// query-string-type query. Suitable for marshalling as JSON and feeding
// to Elasticsearch.
type SourceQuery map[string]interface{}

func (e *QueryEngine) getStartResult(req *SourceRequest) int64 {
	startPage := req.GetStartPage()
	pageSize := req.GetSourcesPerPage()
	if pageSize == 0 {
		return int64(0)
	}
	return startPage * pageSize
}

// getResultCount returns the number of results to return for a
// request.
func (e *QueryEngine) getResultCount(req *SourceRequest) int64 {
	pageSize := req.GetSourcesPerPage()
	if pageSize <= 0 {
		return e.nSources
	}
	return pageSize
}

// BuildQuery takes a SourceRequest and turns it into a multi-level
// map suitable for marshalling to JSON and sending to Elasticsearch.
func (e *QueryEngine) BuildQuery(origin string, req *SourceRequest) (SourceQuery, error) {
	fromResult := e.getStartResult(req)
	resultCount := e.getResultCount(req)
	// First, we check if the query_string field is set; if it is,
	// we can ignore the rest of the request.
	qs := req.GetQueryString()
	if qs != "" {
		query := map[string]interface{}{
			"query": map[string]interface{} {
				"query_string" : map[string]interface{} {
					"query" : qs,
				},
			},
			"from": fromResult,
			"size": resultCount,
		}
		return SourceQuery(query), nil
	}
	tags := req.GetTags()
	tagQueries := make([]*search.QueryDsl, 0)
	for _, tag := range tags {
		q, err := e.buildTagQuery(tag)
		if q != nil && err == nil {
			tagQueries = append(tagQueries, q)
		}
	}
	originTag := e.buildOriginQuery(origin)
	tagQueries = append(tagQueries, originTag)
	if len(tagQueries) == 0 {
		return nil, errors.New("No valid query strings found.")
	}
	query := map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": tagQueries,
			},
		},
		"from": fromResult,
		"size": resultCount,
	}
	return SourceQuery(query), nil
}

// runSourceRequest takes a request object and returns an elastigo-type
// (i.e., intermediate) result.
func (e *QueryEngine) runSourceRequest(origin string, req *SourceRequest) (*es.SearchResult, error) {
	q, err := e.BuildQuery(origin, req)
	if err != nil {
		return nil, err
	}
	var args map[string]interface{}
	res, err := es.SearchRequest(e.indexName, e.dataType, args, q)
	return &res, err
}

// GetSources takes a request object and returns the DataSourceBurst of
// the sources it gets back from Elasticsearch. If error is not nil,
// then a valid DataSourceBurst will still be returned, with the Error
// field set.
func (e *QueryEngine) GetSources(origin string, req *SourceRequest) (*DataSourceBurst, error) {
	res, err := e.runSourceRequest(origin, req)
	if err != nil {
		msg := fmt.Sprintf("Request error: %v", err)
		burst := new(DataSourceBurst)
		burst.Error = &msg
		return burst, err
	}
	sources := make([]*DataSource, len(res.Hits.Hits))
	for i, hit := range res.Hits.Hits {
		source := new(ElasticsearchSource)
		err = json.Unmarshal(*hit.Source, source)
		if err != nil {
			msg := fmt.Sprintf("Response decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, err
		}
		sources[i] = source.Unmarshal()
	}
	burst := BuildSourceBurst(sources)
	return burst, nil
}

// FmtResult returns a string from a SearchResult by interpreting it in
// the most naive manner possible. For debugging.
func FmtResult(result *es.SearchResult) []string {
	results := make([]string, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		results[i] = string((*hit.Source)[:])
	}
	return results
}

// GetSourceCount returns the number of sources we currently think exist
// in the index.
func (e *QueryEngine) GetSourceCount() int64 {
	return e.nSources
}

// updateSourceCount updates our running total of documents-in-index
// (by asking Elasticsearch).
func (e *QueryEngine) updateSourceCount() error {
	var args map[string]interface{}
	resp, err := es.Count(e.indexName, e.dataType, args)
	e.nSources = int64(resp.Count)
	return err
}

// updateForever updates the source counter on a regular basis.
func (e *QueryEngine) updateForever() {
	for true {
		time.Sleep(e.updateInterval)
		err := e.updateSourceCount()
		if err != nil {
			log.Printf("Error updating source count: %v", err)
		}
	}
}
<commit_msg>Refactor a bit<commit_after>package chevalier

import (
	"github.com/mattbaird/elastigo/api"
	es "github.com/mattbaird/elastigo/core"
	"github.com/mattbaird/elastigo/search"

	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"
)

// QueryEngine presents an interface for running queries for sources
// against Elasticsearch.
type QueryEngine struct {
	indexName      string
	metaIndex      string
	originType     string
	dataType       string
	nSources       int64
	updateInterval time.Duration
}

// NewQueryEngine initializes a QueryEngine with the supplied
// Elasticsearch metadata. indexName and dataType can be anything as
// long as they're consistent.
func NewQueryEngine(host, indexName, dataType, metaIndex string) *QueryEngine {
	e := new(QueryEngine)
	e.indexName = indexName
	e.metaIndex = metaIndex
	e.dataType = dataType
	e.originType = "chevalier_origin"
	api.Domain = host
	e.updateSourceCount()
	e.updateInterval = time.Second * 10
	go e.updateForever()
	return e
}

// sanitizeField takes an input field specifier (from a
// SourceRequest_Tag) and munges it so Elasticsearch likes it more. In
// particular, it makes single-wildcard queries work correctly.
//
// FIXME: mid-field wildcards still aren't working.
func (e *QueryEngine) sanitizeField(f string) string {
	f = strings.TrimSpace(f)
	if f == "*" {
		return fmt.Sprintf("%s._all", e.dataType)
	}
	f = fmt.Sprintf("%s.%s", e.dataType, f)
	return f
}

// sanitizeTag escapes assorted Elasticsearch wildcards.
//
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters
func (e *QueryEngine) sanitizeTag(field, value string) (string, string) {
	// * is normally in this list, but is not included here because
	// we want it to act as a wildcard.
	// Also, this can be made a lot faster.
	reservedChars := `\ + - && || ! ( ) { } [ ] ^ " ~ ? : /`
	for _, char := range strings.Split(reservedChars, " ") {
		escapedChar := fmt.Sprintf(`\%s`, char)
		field = strings.Replace(field, char, escapedChar, -1)
		value = strings.Replace(value, char, escapedChar, -1)
	}
	field = e.sanitizeField(field)
	return field, value
}

// buildTagQuery constructs an elastigo query object (search.QueryDsl)
// from a SourceRequest_Tag, designed to be plugged into a
// query-string-type[0] query later on. Returns error on empty query.
//
// [0]: http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html
func (e *QueryEngine) buildTagQuery(tag *SourceRequest_Tag) (*search.QueryDsl, error) {
	field, value := e.sanitizeTag(*tag.Field, *tag.Value)
	// Don't bother running empty queries.
	if value == "" {
		return nil, errors.New("empty query string")
	}
	qs := new(search.QueryString)
	qs.Fields = make([]string, 0)
	qs.Fields = append(qs.Fields, field)
	qs.Query = value
	q := search.Query().Qs(qs)
	return q, nil
}

// buildOriginQuery returns the origin-matching part of the complete
// query.
func (e *QueryEngine) buildOriginQuery(origin string) *search.QueryDsl {
	qs := new(search.QueryString)
	qs.Fields = make([]string, 1)
	qs.Fields[0] = "Origin"
	qs.Query = origin
	q := search.Query().Qs(qs)
	return q
}

// SourceQuery is a multi-level map type representing an Elasticsearch
// query-string-type query. Suitable for marshalling as JSON and feeding
// to Elasticsearch.
type SourceQuery map[string]interface{}

func (e *QueryEngine) getStartResult(req *SourceRequest) int64 {
	startPage := req.GetStartPage()
	pageSize := req.GetSourcesPerPage()
	if pageSize == 0 {
		return int64(0)
	}
	return startPage * pageSize
}

// getResultCount returns the number of results to return for a
// request.
func (e *QueryEngine) getResultCount(req *SourceRequest) int64 {
	pageSize := req.GetSourcesPerPage()
	if pageSize <= 0 {
		return e.nSources
	}
	return pageSize
}

// BuildQuery takes a SourceRequest and turns it into a multi-level
// map suitable for marshalling to JSON and sending to Elasticsearch.
func (e *QueryEngine) BuildQuery(origin string, req *SourceRequest) (SourceQuery, error) {
	fromResult := e.getStartResult(req)
	resultCount := e.getResultCount(req)
	// First, we check if the query_string field is set; if it is,
	// we can ignore the rest of the request.
	qs := req.GetQueryString()
	if qs != "" {
		query := map[string]interface{}{
			"query": map[string]interface{} {
				"query_string" : map[string]interface{} {
					"query" : qs,
				},
			},
			"from": fromResult,
			"size": resultCount,
		}
		return SourceQuery(query), nil
	}
	tags := req.GetTags()
	tagQueries := make([]*search.QueryDsl, 0)
	for _, tag := range tags {
		q, err := e.buildTagQuery(tag)
		if q != nil && err == nil {
			tagQueries = append(tagQueries, q)
		}
	}
	originTag := e.buildOriginQuery(origin)
	tagQueries = append(tagQueries, originTag)
	if len(tagQueries) == 0 {
		return nil, errors.New("No valid query strings found.")
	}
	query := map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": tagQueries,
			},
		},
		"from": fromResult,
		"size": resultCount,
	}
	return SourceQuery(query), nil
}

// runSourceRequest takes a request object and returns an elastigo-type
// (i.e., intermediate) result.
func (e *QueryEngine) runSourceRequest(origin string, req *SourceRequest) (*es.SearchResult, error) {
	q, err := e.BuildQuery(origin, req)
	if err != nil {
		return nil, err
	}
	var args map[string]interface{}
	res, err := es.SearchRequest(e.indexName, e.dataType, args, q)
	return &res, err
}

// GetSources takes a request object and returns the DataSourceBurst of
// the sources it gets back from Elasticsearch. If error is not nil,
// then a valid DataSourceBurst will still be returned, with the Error
// field set.
func (e *QueryEngine) GetSources(origin string, req *SourceRequest) (*DataSourceBurst, error) {
	res, err := e.runSourceRequest(origin, req)
	if err != nil {
		msg := fmt.Sprintf("Request error: %v", err)
		burst := new(DataSourceBurst)
		burst.Error = &msg
		return burst, err
	}
	sources := make([]*DataSource, len(res.Hits.Hits))
	for i, hit := range res.Hits.Hits {
		source := new(ElasticsearchSource)
		err = json.Unmarshal(*hit.Source, source)
		if err != nil {
			msg := fmt.Sprintf("Response decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, err
		}
		sources[i] = source.Unmarshal()
	}
	burst := BuildSourceBurst(sources)
	return burst, nil
}

// FmtResult returns a string from a SearchResult by interpreting it in
// the most naive manner possible. For debugging.
func FmtResult(result *es.SearchResult) []string {
	results := make([]string, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		results[i] = string((*hit.Source)[:])
	}
	return results
}

// GetSourceCount returns the number of sources we currently think exist
// in the index.
func (e *QueryEngine) GetSourceCount() int64 {
	return e.nSources
}

// updateSourceCount updates our running total of documents-in-index
// (by asking Elasticsearch).
func (e *QueryEngine) updateSourceCount() error {
	var args map[string]interface{}
	resp, err := es.Count(e.indexName, e.dataType, args)
	e.nSources = int64(resp.Count)
	return err
}

// updateForever updates the source counter on a regular basis.
func (e *QueryEngine) updateForever() {
	for {
		time.Sleep(e.updateInterval)
		err := e.updateSourceCount()
		if err != nil {
			log.Printf("Error updating source count: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package uniform

import (
	"bytes"
	"encoding/xml"
)

type SelectChild interface{}

type Select struct {
	Options []SelectChild
}

type OptGroup struct {
	Label    string
	Disabled bool
	Options  []Option
}

type Option struct {
	Value    string
	Text     string
	Selected bool
	Disabled bool
}

func (s Select) String() (string, error) {
	buffer := bytes.NewBuffer(nil)
	encoder := errEncoder{encoder: xml.NewEncoder(buffer)}

	encoder.encodeToken(xml.StartElement{
		Name: xml.Name{Local: "select"},
	})

	for _, child := range s.Options {
		switch child.(type) {
		case OptGroup:
			encodeOptGroup(encoder, child.(OptGroup))
		case Option:
			encodeOption(encoder, child.(Option))
		}
	}

	encoder.encodeToken(xml.EndElement{Name: xml.Name{Local: "select"}})
	encoder.flush()

	if encoder.err != nil {
		return "", encoder.err
	}

	return buffer.String(), nil
}

func encodeOptGroup(encoder errEncoder, optGroup OptGroup) {
	optGroupAttrs := make(attributes, 0)
	optGroupAttrs.addString("label", optGroup.Label)
	optGroupAttrs.addBoolean("disabled", optGroup.Disabled)

	encoder.encodeToken(xml.StartElement{
		Name: xml.Name{Local: "optgroup"},
		Attr: optGroupAttrs.xmlAttr(),
	})

	for _, option := range optGroup.Options {
		encodeOption(encoder, option)
	}

	encoder.encodeToken(xml.EndElement{Name: xml.Name{Local: "optgroup"}})
}

func encodeOption(encoder errEncoder, option Option) {
	optionAttrs := make(attributes, 0)
	optionAttrs.addString("value", option.Value)
	optionAttrs.addBoolean("selected", option.Selected)
	optionAttrs.addBoolean("disabled", option.Disabled)

	encoder.encodeToken(xml.StartElement{
		Name: xml.Name{Local: "option"},
		Attr: optionAttrs.xmlAttr(),
	})

	if option.Text != "" {
		encoder.encodeToken(xml.CharData(option.Text))
	}

	encoder.encodeToken(xml.EndElement{Name: xml.Name{Local: "option"}})
}
<commit_msg>Add boilerplate `HTML` method to Select<commit_after>package uniform

import (
	"bytes"
	"encoding/xml"
	"html/template"
)

type SelectChild interface{}

type Select struct {
	Options []SelectChild
}

type OptGroup struct {
	Label    string
	Disabled bool
	Options  []Option
}

type Option struct {
	Value    string
	Text     string
	Selected bool
	Disabled bool
}

func (s Select) String() (string, error) {
	buffer := bytes.NewBuffer(nil)
	encoder := errEncoder{encoder: xml.NewEncoder(buffer)}

	encoder.encodeToken(xml.StartElement{
		Name: xml.Name{Local: "select"},
	})

	for _, child := range s.Options {
		switch child.(type) {
		case OptGroup:
			encodeOptGroup(encoder, child.(OptGroup))
		case Option:
			encodeOption(encoder, child.(Option))
		}
	}

	encoder.encodeToken(xml.EndElement{Name: xml.Name{Local: "select"}})
	encoder.flush()

	if encoder.err != nil {
		return "", encoder.err
	}

	return buffer.String(), nil
}

func (t Select) HTML() (template.HTML, error) {
	markup, err := t.String()
	return template.HTML(markup), err
}

func encodeOptGroup(encoder errEncoder, optGroup OptGroup) {
	optGroupAttrs := make(attributes, 0)
	optGroupAttrs.addString("label", optGroup.Label)
	optGroupAttrs.addBoolean("disabled", optGroup.Disabled)

	encoder.encodeToken(xml.StartElement{
		Name: xml.Name{Local: "optgroup"},
		Attr: optGroupAttrs.xmlAttr(),
	})

	for _, option := range optGroup.Options {
		encodeOption(encoder, option)
	}

	encoder.encodeToken(xml.EndElement{Name: xml.Name{Local: "optgroup"}})
}

func encodeOption(encoder errEncoder, option Option) {
	optionAttrs := make(attributes, 0)
	optionAttrs.addString("value", option.Value)
	optionAttrs.addBoolean("selected", option.Selected)
	optionAttrs.addBoolean("disabled", option.Disabled)

	encoder.encodeToken(xml.StartElement{
		Name: xml.Name{Local: "option"},
		Attr: optionAttrs.xmlAttr(),
	})

	if option.Text != "" {
		encoder.encodeToken(xml.CharData(option.Text))
	}

	encoder.encodeToken(xml.EndElement{Name: xml.Name{Local: "option"}})
}
<|endoftext|>