<commit_before>/*
   DNS-over-HTTPS
   Copyright (C) 2017-2018 Star Brilliant <m13253@hotmail.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 (
	"context"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/http/cookiejar"
	"strings"
	"sync"
	"time"

	"github.com/m13253/dns-over-https/json-dns"
	"github.com/miekg/dns"
	"golang.org/x/net/http2"
)

type Client struct {
	conf                 *config
	bootstrap            []string
	udpServers           []*dns.Server
	tcpServers           []*dns.Server
	bootstrapResolver    *net.Resolver
	cookieJar            http.CookieJar
	httpClientMux        *sync.RWMutex
	httpTransport        *http.Transport
	httpClient           *http.Client
	httpClientLastCreate time.Time
}

type DNSRequest struct {
	response          *http.Response
	reply             *dns.Msg
	udpSize           uint16
	ednsClientAddress net.IP
	ednsClientNetmask uint8
	currentUpstream   string
	err               error
}

func NewClient(conf *config) (c *Client, err error) {
	c = &Client{
		conf: conf,
	}

	udpHandler := dns.HandlerFunc(c.udpHandlerFunc)
	tcpHandler := dns.HandlerFunc(c.tcpHandlerFunc)
	for _, addr := range conf.Listen {
		c.udpServers = append(c.udpServers, &dns.Server{
			Addr:    addr,
			Net:     "udp",
			Handler: udpHandler,
			UDPSize: dns.DefaultMsgSize,
		})
		c.tcpServers = append(c.tcpServers, &dns.Server{
			Addr:    addr,
			Net:     "tcp",
			Handler: tcpHandler,
		})
	}
	c.bootstrapResolver = net.DefaultResolver
	if len(conf.Bootstrap) != 0 {
		c.bootstrap = make([]string, len(conf.Bootstrap))
		for i, bootstrap := range conf.Bootstrap {
			bootstrapAddr, err := net.ResolveUDPAddr("udp", bootstrap)
			if err != nil {
				bootstrapAddr, err = net.ResolveUDPAddr("udp", "["+bootstrap+"]:53")
			}
			if err != nil {
				return nil, err
			}
			c.bootstrap[i] = bootstrapAddr.String()
		}
		c.bootstrapResolver = &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				var d net.Dialer
				numServers := len(c.bootstrap)
				bootstrap := c.bootstrap[rand.Intn(numServers)]
				conn, err := d.DialContext(ctx, network, bootstrap)
				return conn, err
			},
		}
	}
	// Most CDNs require Cookie support to prevent DDoS attack.
	// Disabling Cookie does not effectively prevent tracking,
	// so I will leave it on to make anti-DDoS services happy.
	if !c.conf.NoCookies {
		c.cookieJar, err = cookiejar.New(nil)
		if err != nil {
			return nil, err
		}
	} else {
		c.cookieJar = nil;
	}

	c.httpClientMux = new(sync.RWMutex)
	err = c.newHTTPClient()
	if err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Client) newHTTPClient() error {
	c.httpClientMux.Lock()
	defer c.httpClientMux.Unlock()
	if !c.httpClientLastCreate.IsZero() && time.Now().Sub(c.httpClientLastCreate) < time.Duration(c.conf.Timeout)*time.Second {
		return nil
	}
	if c.httpTransport != nil {
		c.httpTransport.CloseIdleConnections()
	}
	dialer := &net.Dialer{
		Timeout:   time.Duration(c.conf.Timeout) * time.Second,
		KeepAlive: 30 * time.Second,
		DualStack: true,
		Resolver:  c.bootstrapResolver,
	}
	c.httpTransport = &http.Transport{
		DialContext:           dialer.DialContext,
		ExpectContinueTimeout: 1 * time.Second,
		IdleConnTimeout:       90 * time.Second,
		MaxIdleConns:          100,
		MaxIdleConnsPerHost:   10,
		Proxy:                 http.ProxyFromEnvironment,
		ResponseHeaderTimeout: time.Duration(c.conf.Timeout) * time.Second,
		TLSHandshakeTimeout:   time.Duration(c.conf.Timeout) * time.Second,
	}
	if c.conf.NoIPv6 {
		c.httpTransport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
			if strings.HasPrefix(network, "tcp") {
				network = "tcp4"
			}
			return dialer.DialContext(ctx, network, address)
		}
	}
	err := http2.ConfigureTransport(c.httpTransport)
	if err != nil {
		return err
	}
	c.httpClient = &http.Client{
		Transport: c.httpTransport,
		Jar:       c.cookieJar,
	}
	c.httpClientLastCreate = time.Now()
	return nil
}

func (c *Client) Start() error {
	results := make(chan error, len(c.udpServers)+len(c.tcpServers))
	for _, srv := range append(c.udpServers, c.tcpServers...) {
		go func(srv *dns.Server) {
			err := srv.ListenAndServe()
			if err != nil {
				log.Println(err)
			}
			results <- err
		}(srv)
	}

	for i := 0; i < cap(results); i++ {
		err := <-results
		if err != nil {
			return err
		}
	}
	close(results)
	return nil
}

func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
	if r.Response == true {
		log.Println("Received a response packet")
		return
	}

	requestType := ""
	if len(c.conf.UpstreamIETF) == 0 {
		requestType = "application/dns-json"
	} else if len(c.conf.UpstreamGoogle) == 0 {
		requestType = "application/dns-message"
	} else {
		numServers := len(c.conf.UpstreamGoogle) + len(c.conf.UpstreamIETF)
		random := rand.Intn(numServers)
		if random < len(c.conf.UpstreamGoogle) {
			requestType = "application/dns-json"
		} else {
			requestType = "application/dns-message"
		}
	}

	var req *DNSRequest
	if requestType == "application/dns-json" {
		req = c.generateRequestGoogle(w, r, isTCP)
	} else if requestType == "application/dns-message" {
		req = c.generateRequestIETF(w, r, isTCP)
	} else {
		panic("Unknown request Content-Type")
	}

	if req.response != nil {
		for _, header := range c.conf.DebugHTTPHeaders {
			if value := req.response.Header.Get(header); value != "" {
				log.Printf("%s: %s\n", header, value)
			}
		}
	}
	if req.err != nil {
		return
	}

	contentType := ""
	candidateType := strings.SplitN(req.response.Header.Get("Content-Type"), ";", 2)[0]
	if candidateType == "application/json" {
		contentType = "application/json"
	} else if candidateType == "application/dns-message" {
		contentType = "application/dns-message"
	} else if candidateType == "application/dns-udpwireformat" {
		contentType = "application/dns-message"
	} else {
		if requestType == "application/dns-json" {
			contentType = "application/json"
		} else if requestType == "application/dns-message" {
			contentType = "application/dns-message"
		}
	}

	if contentType == "application/json" {
		c.parseResponseGoogle(w, r, isTCP, req)
	} else if contentType == "application/dns-message" {
		c.parseResponseIETF(w, r, isTCP, req)
	} else {
		panic("Unknown response Content-Type")
	}
}

func (c *Client) udpHandlerFunc(w dns.ResponseWriter, r *dns.Msg) {
	c.handlerFunc(w, r, false)
}

func (c *Client) tcpHandlerFunc(w dns.ResponseWriter, r *dns.Msg) {
	c.handlerFunc(w, r, true)
}

var (
	ipv4Mask24 = net.IPMask{255, 255, 255, 0}
	ipv6Mask56 = net.CIDRMask(56, 128)
)

func (c *Client) findClientIP(w dns.ResponseWriter, r *dns.Msg) (ednsClientAddress net.IP, ednsClientNetmask uint8) {
	ednsClientNetmask = 255
	if c.conf.NoECS {
		return net.IPv4(0, 0, 0, 0), 0
	}
	if opt := r.IsEdns0(); opt != nil {
		for _, option := range opt.Option {
			if option.Option() == dns.EDNS0SUBNET {
				edns0Subnet := option.(*dns.EDNS0_SUBNET)
				ednsClientAddress = edns0Subnet.Address
				ednsClientNetmask = edns0Subnet.SourceNetmask
				return
			}
		}
	}
	remoteAddr, err := net.ResolveUDPAddr("udp", w.RemoteAddr().String())
	if err != nil {
		return
	}
	if ip := remoteAddr.IP; jsonDNS.IsGlobalIP(ip) {
		if ipv4 := ip.To4(); ipv4 != nil {
			ednsClientAddress = ipv4.Mask(ipv4Mask24)
			ednsClientNetmask = 24
		} else {
			ednsClientAddress = ip.Mask(ipv6Mask56)
			ednsClientNetmask = 56
		}
	}
	return
}
<commit_msg>Fix HTTP stream leaking problem<commit_after>/*
   DNS-over-HTTPS
   Copyright (C) 2017-2018 Star Brilliant <m13253@hotmail.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 (
	"context"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/http/cookiejar"
	"strings"
	"sync"
	"time"

	"github.com/m13253/dns-over-https/json-dns"
	"github.com/miekg/dns"
	"golang.org/x/net/http2"
)

type Client struct {
	conf                 *config
	bootstrap            []string
	udpServers           []*dns.Server
	tcpServers           []*dns.Server
	bootstrapResolver    *net.Resolver
	cookieJar            http.CookieJar
	httpClientMux        *sync.RWMutex
	httpTransport        *http.Transport
	httpClient           *http.Client
	httpClientLastCreate time.Time
}

type DNSRequest struct {
	response          *http.Response
	reply             *dns.Msg
	udpSize           uint16
	ednsClientAddress net.IP
	ednsClientNetmask uint8
	currentUpstream   string
	err               error
}

func NewClient(conf *config) (c *Client, err error) {
	c = &Client{
		conf: conf,
	}

	udpHandler := dns.HandlerFunc(c.udpHandlerFunc)
	tcpHandler := dns.HandlerFunc(c.tcpHandlerFunc)
	for _, addr := range conf.Listen {
		c.udpServers = append(c.udpServers, &dns.Server{
			Addr:    addr,
			Net:     "udp",
			Handler: udpHandler,
			UDPSize: dns.DefaultMsgSize,
		})
		c.tcpServers = append(c.tcpServers, &dns.Server{
			Addr:    addr,
			Net:     "tcp",
			Handler: tcpHandler,
		})
	}
	c.bootstrapResolver = net.DefaultResolver
	if len(conf.Bootstrap) != 0 {
		c.bootstrap = make([]string, len(conf.Bootstrap))
		for i, bootstrap := range conf.Bootstrap {
			bootstrapAddr, err := net.ResolveUDPAddr("udp", bootstrap)
			if err != nil {
				bootstrapAddr, err = net.ResolveUDPAddr("udp", "["+bootstrap+"]:53")
			}
			if err != nil {
				return nil, err
			}
			c.bootstrap[i] = bootstrapAddr.String()
		}
		c.bootstrapResolver = &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				var d net.Dialer
				numServers := len(c.bootstrap)
				bootstrap := c.bootstrap[rand.Intn(numServers)]
				conn, err := d.DialContext(ctx, network, bootstrap)
				return conn, err
			},
		}
	}
	// Most CDNs require Cookie support to prevent DDoS attack.
	// Disabling Cookie does not effectively prevent tracking,
	// so I will leave it on to make anti-DDoS services happy.
	if !c.conf.NoCookies {
		c.cookieJar, err = cookiejar.New(nil)
		if err != nil {
			return nil, err
		}
	} else {
		c.cookieJar = nil;
	}

	c.httpClientMux = new(sync.RWMutex)
	err = c.newHTTPClient()
	if err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Client) newHTTPClient() error {
	c.httpClientMux.Lock()
	defer c.httpClientMux.Unlock()
	if !c.httpClientLastCreate.IsZero() && time.Now().Sub(c.httpClientLastCreate) < time.Duration(c.conf.Timeout)*time.Second {
		return nil
	}
	if c.httpTransport != nil {
		c.httpTransport.CloseIdleConnections()
	}
	dialer := &net.Dialer{
		Timeout:   time.Duration(c.conf.Timeout) * time.Second,
		KeepAlive: 30 * time.Second,
		DualStack: true,
		Resolver:  c.bootstrapResolver,
	}
	c.httpTransport = &http.Transport{
		DialContext:           dialer.DialContext,
		ExpectContinueTimeout: 1 * time.Second,
		IdleConnTimeout:       90 * time.Second,
		MaxIdleConns:          100,
		MaxIdleConnsPerHost:   10,
		Proxy:                 http.ProxyFromEnvironment,
		ResponseHeaderTimeout: time.Duration(c.conf.Timeout) * time.Second,
		TLSHandshakeTimeout:   time.Duration(c.conf.Timeout) * time.Second,
	}
	if c.conf.NoIPv6 {
		c.httpTransport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
			if strings.HasPrefix(network, "tcp") {
				network = "tcp4"
			}
			return dialer.DialContext(ctx, network, address)
		}
	}
	err := http2.ConfigureTransport(c.httpTransport)
	if err != nil {
		return err
	}
	c.httpClient = &http.Client{
		Transport: c.httpTransport,
		Jar:       c.cookieJar,
	}
	c.httpClientLastCreate = time.Now()
	return nil
}

func (c *Client) Start() error {
	results := make(chan error, len(c.udpServers)+len(c.tcpServers))
	for _, srv := range append(c.udpServers, c.tcpServers...) {
		go func(srv *dns.Server) {
			err := srv.ListenAndServe()
			if err != nil {
				log.Println(err)
			}
			results <- err
		}(srv)
	}

	for i := 0; i < cap(results); i++ {
		err := <-results
		if err != nil {
			return err
		}
	}
	close(results)
	return nil
}

func (c *Client) handlerFunc(w dns.ResponseWriter, r *dns.Msg, isTCP bool) {
	if r.Response == true {
		log.Println("Received a response packet")
		return
	}

	requestType := ""
	if len(c.conf.UpstreamIETF) == 0 {
		requestType = "application/dns-json"
	} else if len(c.conf.UpstreamGoogle) == 0 {
		requestType = "application/dns-message"
	} else {
		numServers := len(c.conf.UpstreamGoogle) + len(c.conf.UpstreamIETF)
		random := rand.Intn(numServers)
		if random < len(c.conf.UpstreamGoogle) {
			requestType = "application/dns-json"
		} else {
			requestType = "application/dns-message"
		}
	}

	var req *DNSRequest
	if requestType == "application/dns-json" {
		req = c.generateRequestGoogle(w, r, isTCP)
	} else if requestType == "application/dns-message" {
		req = c.generateRequestIETF(w, r, isTCP)
	} else {
		panic("Unknown request Content-Type")
	}

	if req.response != nil {
		defer req.response.Body.Close()
		for _, header := range c.conf.DebugHTTPHeaders {
			if value := req.response.Header.Get(header); value != "" {
				log.Printf("%s: %s\n", header, value)
			}
		}
	}
	if req.err != nil {
		return
	}

	contentType := ""
	candidateType := strings.SplitN(req.response.Header.Get("Content-Type"), ";", 2)[0]
	if candidateType == "application/json" {
		contentType = "application/json"
	} else if candidateType == "application/dns-message" {
		contentType = "application/dns-message"
	} else if candidateType == "application/dns-udpwireformat" {
		contentType = "application/dns-message"
	} else {
		if requestType == "application/dns-json" {
			contentType = "application/json"
		} else if requestType == "application/dns-message" {
			contentType = "application/dns-message"
		}
	}

	if contentType == "application/json" {
		c.parseResponseGoogle(w, r, isTCP, req)
	} else if contentType == "application/dns-message" {
		c.parseResponseIETF(w, r, isTCP, req)
	} else {
		panic("Unknown response Content-Type")
	}
}

func (c *Client) udpHandlerFunc(w dns.ResponseWriter, r *dns.Msg) {
	c.handlerFunc(w, r, false)
}

func (c *Client) tcpHandlerFunc(w dns.ResponseWriter, r *dns.Msg) {
	c.handlerFunc(w, r, true)
}

var (
	ipv4Mask24 = net.IPMask{255, 255, 255, 0}
	ipv6Mask56 = net.CIDRMask(56, 128)
)

func (c *Client) findClientIP(w dns.ResponseWriter, r *dns.Msg) (ednsClientAddress net.IP, ednsClientNetmask uint8) {
	ednsClientNetmask = 255
	if c.conf.NoECS {
		return net.IPv4(0, 0, 0, 0), 0
	}
	if opt := r.IsEdns0(); opt != nil {
		for _, option := range opt.Option {
			if option.Option() == dns.EDNS0SUBNET {
				edns0Subnet := option.(*dns.EDNS0_SUBNET)
				ednsClientAddress = edns0Subnet.Address
				ednsClientNetmask = edns0Subnet.SourceNetmask
				return
			}
		}
	}
	remoteAddr, err := net.ResolveUDPAddr("udp", w.RemoteAddr().String())
	if err != nil {
		return
	}
	if ip := remoteAddr.IP; jsonDNS.IsGlobalIP(ip) {
		if ipv4 := ip.To4(); ipv4 != nil {
			ednsClientAddress = ipv4.Mask(ipv4Mask24)
			ednsClientNetmask = 24
		} else {
			ednsClientAddress = ip.Mask(ipv6Mask56)
			ednsClientNetmask = 56
		}
	}
	return
}
<|endoftext|>


<commit_before>package integration

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/big"
	"os/exec"
	"s3cli/config"
	"time"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

const alphanum = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

// GenerateRandomString generates a random string of len 25
func GenerateRandomString() string {
	size := 25
	randBytes := make([]byte, size)
	for i := range randBytes {
		randInt, err := rand.Int(rand.Reader, big.NewInt(int64(len(alphanum))))
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		randBytes[i] = alphanum[randInt.Uint64()]
	}
	return string(randBytes)
}

// MakeConfigFile creates a config file from a S3Cli config struct
func MakeConfigFile(cfg *config.S3Cli) string {
	cfgBytes, err := json.Marshal(cfg)
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	tmpFile, err := ioutil.TempFile("", "s3cli-test")
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	_, err = tmpFile.Write(cfgBytes)
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	err = tmpFile.Close()
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	return tmpFile.Name()
}

// MakeContentFile creates a temporary file with content to upload to S3
func MakeContentFile(content string) string {
	tmpFile, err := ioutil.TempFile("", "s3cli-test-content")
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	_, err = tmpFile.Write([]byte(content))
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	err = tmpFile.Close()
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	return tmpFile.Name()
}

// RunS3CLI runs the s3cli and outputs the session after waiting for it to finish
func RunS3CLI(s3CLIPath string, configPath string, subcommand string, args ...string) (*gexec.Session, error) {
	cmdArgs := []string{
		fmt.Sprintf("--config %s", configPath),
		subcommand,
	}
	cmdArgs = append(cmdArgs, args...)
	command := exec.Command(s3CLIPath, cmdArgs...)
	gexecSession, err := gexec.Start(command, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	if err != nil {
		return nil, err
	}
	gexecSession.Wait(1 * time.Minute)
	return gexecSession, nil
}
<commit_msg>use the correct flag to specify path to configuration for tests.<commit_after>package integration

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/big"
	"os/exec"
	"s3cli/config"
	"time"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

const alphanum = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

// GenerateRandomString generates a random string of len 25
func GenerateRandomString() string {
	size := 25
	randBytes := make([]byte, size)
	for i := range randBytes {
		randInt, err := rand.Int(rand.Reader, big.NewInt(int64(len(alphanum))))
		gomega.Expect(err).ToNot(gomega.HaveOccurred())
		randBytes[i] = alphanum[randInt.Uint64()]
	}
	return string(randBytes)
}

// MakeConfigFile creates a config file from a S3Cli config struct
func MakeConfigFile(cfg *config.S3Cli) string {
	cfgBytes, err := json.Marshal(cfg)
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	tmpFile, err := ioutil.TempFile("", "s3cli-test")
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	_, err = tmpFile.Write(cfgBytes)
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	err = tmpFile.Close()
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	return tmpFile.Name()
}

// MakeContentFile creates a temporary file with content to upload to S3
func MakeContentFile(content string) string {
	tmpFile, err := ioutil.TempFile("", "s3cli-test-content")
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	_, err = tmpFile.Write([]byte(content))
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	err = tmpFile.Close()
	gomega.Expect(err).ToNot(gomega.HaveOccurred())
	return tmpFile.Name()
}

// RunS3CLI runs the s3cli and outputs the session after waiting for it to finish
func RunS3CLI(s3CLIPath string, configPath string, subcommand string, args ...string) (*gexec.Session, error) {
	cmdArgs := []string{
		fmt.Sprintf("-c %s", configPath),
		subcommand,
	}
	cmdArgs = append(cmdArgs, args...)
	command := exec.Command(s3CLIPath, cmdArgs...)
	gexecSession, err := gexec.Start(command, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	if err != nil {
		return nil, err
	}
	gexecSession.Wait(1 * time.Minute)
	return gexecSession, nil
}
<|endoftext|>


<commit_before>package notifiers

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/InVisionApp/go-datadog-api"
	"github.com/InVisionApp/kit-overwatch/notifiers/deps"
	log "github.com/Sirupsen/logrus"
	"strings"
)

const EVENT_TYPE = "kubernetes"
const MD_PREFIX = "%%% \n"
const MD_SUFFIX = "\n %%%"

type NotifyDataDog struct {
	Client *datadog.Client
}

func New(apiKey string, appKey string) *NotifyDataDog {
	client := datadog.NewClient(apiKey, appKey)
	return &NotifyDataDog{
		Client: client,
	}
}

func (ndd *NotifyDataDog) Send(n *deps.Notification) error {

	event := &datadog.Event{
		SourceType: EVENT_TYPE,
		EventType:  EVENT_TYPE,
	}

	switch n.Level {
	case "INFO":
		event.AlertType = "Info"
		event.Priority = "low"
	case "WARN":
		event.AlertType = "Warning"
		event.Priority = "high"
	case "ERROR":
		event.AlertType = "Error"
		event.Priority = "high"
	}

	serviceName := n.Event.ObjectMeta.Name
	splitName := strings.Split(n.Event.ObjectMeta.Name, "-")

	switch n.Event.InvolvedObject.Kind {
	case "Pod":
		// Default Rules to ignore `service-name-[deploynumber-podnumber.number]`
		indexToIgnore := 2

		if len(splitName) > 3 && splitName[len(splitName)-3] == "deployment" {
			// Rules to ignore `service-name-[deployment-deploynumber-podnumber.number]`
			indexToIgnore = 3
		}
		serviceName = strings.Join(splitName[:len(splitName)-indexToIgnore], "-")

	}

	title := fmt.Sprintf("`%s` event for `%s` on `%s`", n.Event.Reason, serviceName, n.Cluster)
	if n.Mention != "" {
		title = fmt.Sprintf("Alerting [%s] concerning %s", n.Mention, title)
	}

	event.Title = title
	event.Tags = []string{
		"team:" + n.Mention,
		"cluster:" + n.Cluster,
		"type:" + n.Event.Type,
		"level:" + n.Level,
		"reason:" + n.Event.Reason,
		"node:" + n.Event.Source.Host,
		"name:" + n.Event.ObjectMeta.Name,
		"namespace:" + n.Event.ObjectMeta.Namespace,
		"component:" + n.Event.Source.Component,
		"count:" + fmt.Sprintf("%d", n.Event.Count),
		"object-kind:" + n.Event.InvolvedObject.Kind,
		"object-name:" + n.Event.InvolvedObject.Name,
		"mentioned:" + n.Mention,
		"service:" + serviceName,
	}

	message := `#### Message Details
	%v
#### Event Details
	%v
#### Involved Object
	%v
`

	mDetails := &messageDetails{
		Cluster: n.Cluster,
		Reason:  n.Event.Reason,
		Type:    n.Event.Type,
		Level:   n.Level,
	}

	eDetails := &eventDetails{
		Node:           n.Event.Source.Host,
		Namespace:      n.Event.ObjectMeta.Namespace,
		Component:      n.Event.Source.Component,
		Count:          fmt.Sprintf("%d", n.Event.Count),
		FirstOccurance: n.Event.FirstTimestamp.Format(time.RFC1123),
		LastOccurance:  n.Event.LastTimestamp.Format(time.RFC1123),
	}

	iObject := &involvedObject{
		Kind: n.Event.InvolvedObject.Kind,
		Name: n.Event.InvolvedObject.Name,
	}

	mDetailsJson, err := json.Marshal(mDetails)
	if err != nil {
		return fmt.Errorf("%s\n", err)
	}
	mDetailsFmt := string(mDetailsJson)

	eDetailsJson, err := json.Marshal(eDetails)
	if err != nil {
		return fmt.Errorf("%s\n", err)
	}
	eDetailsFmt := string(eDetailsJson)

	iObjectJson, err := json.Marshal(iObject)
	if err != nil {
		return fmt.Errorf("%s\n", err)
	}
	iObjectFmt := string(iObjectJson)

	event.Text = MD_PREFIX + fmt.Sprintf(message, mDetailsFmt, eDetailsFmt, iObjectFmt) + MD_SUFFIX

	newEvent, err := ndd.Client.PostEvent(event)
	if err != nil {
		return fmt.Errorf("Post Event error: %s\n", err)
	}

	log.Infof("NotifyDataDog: %s / %s / %s / %s", n.Event.Reason, n.Event.Message, newEvent.Id, newEvent.Title)
	return nil
}

type messageDetails struct {
	Cluster string `json:"cluster,omitempty"`
	Reason  string `json:"reason,omitempty"`
	Type    string `json:"type,omitempty"`
	Level   string `json:"level,omitempty"`
}

type eventDetails struct {
	Node           string `json:"node,omitempty"`
	Namespace      string `json:"namespace,omitempty"`
	Component      string `json:"component,omitempty"`
	Count          string `json:"count,omitempty"`
	FirstOccurance string `json:"first_occurance,omitempty"`
	LastOccurance  string `json:"last_occurance,omitempty"`
}

type involvedObject struct {
	Kind string `json:"kind,omitempty"`
	Name string `json:"name,omitempty"`
}
<commit_msg>Fixing comment to read replicaset (correctly )<commit_after>package notifiers

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/zorkian/go-datadog-api"

	dependencies "github.com/InVisionApp/kit-overwatch/deps"
	"github.com/InVisionApp/kit-overwatch/notifiers/deps"
)

const EVENT_TYPE = "kubernetes"
const MD_PREFIX = "%%% \n"
const MD_SUFFIX = "\n %%%"

type NotifyDataDog struct {
	Client *datadog.Client
}

func New(apiKey string, appKey string) *NotifyDataDog {
	client := datadog.NewClient(apiKey, appKey)
	return &NotifyDataDog{
		Client: client,
	}
}

func (ndd *NotifyDataDog) Send(n *deps.Notification) error {

	event := &datadog.Event{
		SourceType: EVENT_TYPE,
		EventType:  EVENT_TYPE,
	}

	switch n.Level {
	case "INFO":
		event.AlertType = "Info"
		event.Priority = "low"
	case "WARN":
		event.AlertType = "Warning"
		event.Priority = "high"
	case "ERROR":
		event.AlertType = "Error"
		event.Priority = "high"
	}

	serviceName := n.Event.ObjectMeta.Name
	splitName := strings.Split(n.Event.ObjectMeta.Name, "-")

	switch n.Event.InvolvedObject.Kind {
	case "Pod":
		// Default Rules to ignore `service-name-[replicaset-podnumber.number]`
		indexToIgnore := 2

		if len(splitName) > 3 && splitName[len(splitName)-3] == "deployment" {
			// Rules to ignore `service-name-[deployment-replicaset-podnumber.number]`
			indexToIgnore = 3
		}
		serviceName = strings.Join(splitName[:len(splitName)-indexToIgnore], "-")

	}

	title := fmt.Sprintf("`%s` event for `%s` on `%s`", n.Event.Reason, serviceName, n.Cluster)
	if n.Mention != "" {
		title = fmt.Sprintf("Alerting [%s] concerning %s", n.Mention, title)
	}

	event.Title = title
	event.Tags = []string{
		"team:" + n.Mention,
		"cluster:" + n.Cluster,
		"type:" + n.Event.Type,
		"level:" + n.Level,
		"reason:" + n.Event.Reason,
		"node:" + n.Event.Source.Host,
		"name:" + n.Event.ObjectMeta.Name,
		"namespace:" + n.Event.ObjectMeta.Namespace,
		"component:" + n.Event.Source.Component,
		"count:" + fmt.Sprintf("%d", n.Event.Count),
		"object-kind:" + n.Event.InvolvedObject.Kind,
		"object-name:" + n.Event.InvolvedObject.Name,
		"mentioned:" + n.Mention,
		"service:" + serviceName,
	}

	message := `#### Message Details
	%v
#### Event Details
	%v
#### Involved Object
	%v
`

	mDetails := &messageDetails{
		Cluster: n.Cluster,
		Reason:  n.Event.Reason,
		Type:    n.Event.Type,
		Level:   n.Level,
	}

	eDetails := &eventDetails{
		Node:           n.Event.Source.Host,
		Namespace:      n.Event.ObjectMeta.Namespace,
		Component:      n.Event.Source.Component,
		Count:          fmt.Sprintf("%d", n.Event.Count),
		FirstOccurance: n.Event.FirstTimestamp.Format(time.RFC1123),
		LastOccurance:  n.Event.LastTimestamp.Format(time.RFC1123),
	}

	iObject := &involvedObject{
		Kind: n.Event.InvolvedObject.Kind,
		Name: n.Event.InvolvedObject.Name,
	}

	mDetailsJson, err := json.Marshal(mDetails)
	if err != nil {
		return fmt.Errorf("%s\n", err)
	}
	mDetailsFmt := string(mDetailsJson)

	eDetailsJson, err := json.Marshal(eDetails)
	if err != nil {
		return fmt.Errorf("%s\n", err)
	}
	eDetailsFmt := string(eDetailsJson)

	iObjectJson, err := json.Marshal(iObject)
	if err != nil {
		return fmt.Errorf("%s\n", err)
	}
	iObjectFmt := string(iObjectJson)

	event.Text = MD_PREFIX + fmt.Sprintf(message, mDetailsFmt, eDetailsFmt, iObjectFmt) + MD_SUFFIX

	newEvent, err := ndd.Client.PostEvent(event)
	if err != nil {
		return fmt.Errorf("Post Event error: %s\n", err)
	}

	log.Infof("NotifyDataDog: %s / %s / %s / %s", n.Event.Reason, n.Event.Message, newEvent.Id, newEvent.Title)
	return nil
}

type messageDetails struct {
	Cluster string `json:"cluster,omitempty"`
	Reason  string `json:"reason,omitempty"`
	Type    string `json:"type,omitempty"`
	Level   string `json:"level,omitempty"`
}

type eventDetails struct {
	Node           string `json:"node,omitempty"`
	Namespace      string `json:"namespace,omitempty"`
	Component      string `json:"component,omitempty"`
	Count          string `json:"count,omitempty"`
	FirstOccurance string `json:"first_occurance,omitempty"`
	LastOccurance  string `json:"last_occurance,omitempty"`
}

type involvedObject struct {
	Kind string `json:"kind,omitempty"`
	Name string `json:"name,omitempty"`
}
<|endoftext|>


<commit_before>package tchannel

// Copyright (c) 2015 Uber Technologies, Inc.

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

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

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

var (
	errAlreadyListening  = errors.New("channel already listening")
	errInvalidStateForOp = errors.New("channel is in an invalid state for that operation")
)

const (
	ephemeralHostPort = "0.0.0.0:0"
)

// ChannelOptions are used to control parameters on a create a TChannel
type ChannelOptions struct {
	// Default Connection options
	DefaultConnectionOptions ConnectionOptions

	// The name of the process, for logging and reporting to peers
	ProcessName string

	// The logger to use for this channel
	Logger Logger

	// The reporter to use for reporting stats for this channel.
	StatsReporter StatsReporter

	// Trace reporter to use for this channel.
	TraceReporter TraceReporter
}

// ChannelState is the state of a channel.
type ChannelState int

const (
	// ChannelClient is a channel that can be used as a client.
	ChannelClient ChannelState = iota + 1

	// ChannelListening is a channel that is listening for new connnections.
	ChannelListening

	// ChannelStartClose is a channel that has received a Close request.
	// The channel is no longer listening, and all new incoming connections are rejected.
	ChannelStartClose

	// ChannelInboundClosed is a channel that has drained all incoming connections, but may
	// have outgoing connections. All incoming calls and new outgoing calls are rejected.
	ChannelInboundClosed

	// ChannelClosed is a channel that has closed completely.
	ChannelClosed
)

//go:generate stringer -type=ChannelState

// A Channel is a bi-directional connection to the peering and routing network.
// Applications can use a Channel to make service calls to remote peers via
// BeginCall, or to listen for incoming calls from peers.  Applications that
// want to receive requests should call one of Serve or ListenAndServe
// TODO(prashant): Shutdown all subchannels + peers when channel is closed.
type Channel struct {
	log               Logger
	commonStatsTags   map[string]string
	statsReporter     StatsReporter
	traceReporter     TraceReporter
	connectionOptions ConnectionOptions
	handlers          *handlerMap
	peers             *PeerList
	subChannels       *subChannelMap

	// mutable contains all the members of Channel which are mutable.
	mutable struct {
		mut      sync.RWMutex // protects members of the mutable struct.
		state    ChannelState
		peerInfo LocalPeerInfo // May be ephemeral if this is a client only channel
		l        net.Listener  // May be nil if this is a client only channel
		conns    []*Connection
	}
}

// NewChannel creates a new Channel.  The new channel can be used to send outbound requests
// to peers, but will not listen or handling incoming requests until one of ListenAndServe
// or Serve is called. The local service name should be passed to serviceName.
func NewChannel(serviceName string, opts *ChannelOptions) (*Channel, error) {
	if opts == nil {
		opts = &ChannelOptions{}
	}

	logger := opts.Logger
	if logger == nil {
		logger = NullLogger
	}

	traceReporter := opts.TraceReporter
	if traceReporter == nil {
		traceReporter = NullReporter
	}

	processName := opts.ProcessName
	if processName == "" {
		processName = fmt.Sprintf("%s[%d]", filepath.Base(os.Args[0]), os.Getpid())
	}

	statsReporter := opts.StatsReporter
	if statsReporter == nil {
		statsReporter = NullStatsReporter
	}

	ch := &Channel{
		connectionOptions: opts.DefaultConnectionOptions,
		log:               logger.WithFields(LogField{"service", serviceName}),
		statsReporter:     statsReporter,
		traceReporter:     traceReporter,
		handlers:          &handlerMap{},
		subChannels:       &subChannelMap{},
	}
	ch.mutable.peerInfo = LocalPeerInfo{
		PeerInfo: PeerInfo{
			ProcessName: processName,
			HostPort:    ephemeralHostPort,
		},
		ServiceName: serviceName,
	}
	ch.mutable.state = ChannelClient
	ch.peers = newPeerList(ch)
	ch.createCommonStats()
	return ch, nil
}

// Serve serves incoming requests using the provided listener.
// The local peer info is set synchronously, but the actual socket listening is done in
// a separate goroutine.
func (ch *Channel) Serve(l net.Listener) error {
	mutable := &ch.mutable
	mutable.mut.Lock()
	defer mutable.mut.Unlock()

	if mutable.l != nil {
		return errAlreadyListening
	}
	mutable.l = l

	if mutable.state != ChannelClient {
		return errInvalidStateForOp
	}
	mutable.state = ChannelListening

	mutable.peerInfo.HostPort = l.Addr().String()
	peerInfo := mutable.peerInfo
	ch.log.Debugf("%v (%v) listening on %v", peerInfo.ProcessName, peerInfo.ServiceName, peerInfo.HostPort)
	go ch.serve()
	return nil
}

// ListenAndServe listens on the given address and serves incoming requests.
// The port may be 0, in which case the channel will use an OS assigned port
// This method does not block as the handling of connections is done in a goroutine.
func (ch *Channel) ListenAndServe(hostPort string) error {
	mutable := &ch.mutable
	mutable.mut.RLock()

	if mutable.l != nil {
		mutable.mut.RUnlock()
		return errAlreadyListening
	}

	l, err := net.Listen("tcp", hostPort)
	if err != nil {
		mutable.mut.RUnlock()
		return err
	}

	mutable.mut.RUnlock()
	return ch.Serve(l)
}

// Registrar is the base interface for registering handlers on either the base
// Channel or the SubChannel
type Registrar interface {
	// ServiceName returns the service name that this Registrar is for.
	ServiceName() string

	// Register registers a handler for ServiceName and the given operation.
	Register(h Handler, operationName string)

	// Logger returns the logger for this Registrar.
	Logger() Logger
}

// Register registers a handler for a service+operation pair
func (ch *Channel) Register(h Handler, operationName string) {
	ch.handlers.register(h, ch.PeerInfo().ServiceName, operationName)
}

// PeerInfo returns the current peer info for the channel
func (ch *Channel) PeerInfo() LocalPeerInfo {
	ch.mutable.mut.RLock()
	defer ch.mutable.mut.RUnlock()

	return ch.mutable.peerInfo
}

func (ch *Channel) createCommonStats() {
	ch.commonStatsTags = map[string]string{
		"app":     ch.mutable.peerInfo.ProcessName,
		"service": ch.mutable.peerInfo.ServiceName,
	}
	host, err := os.Hostname()
	if err != nil {
		ch.log.Infof("channel failed to get host: %v", err)
		return
	}
	ch.commonStatsTags["host"] = host
	// TODO(prashant): Allow user to pass extra tags (such as cluster, version).
}

// GetSubChannel returns a SubChannel for the given service name. If the subchannel does not
// exist, it is created.
func (ch *Channel) GetSubChannel(serviceName string) *SubChannel {
	return ch.subChannels.getOrAdd(serviceName, ch)
}

// Peers returns the PeerList for the channel.
func (ch *Channel) Peers() *PeerList {
	return ch.peers
}

// BeginCall starts a new call to a remote peer, returning an OutboundCall that can
// be used to write the arguments of the call.
func (ch *Channel) BeginCall(ctx context.Context, hostPort, serviceName, operationName string, callOptions *CallOptions) (*OutboundCall, error) {
	p := ch.peers.GetOrAdd(hostPort)
	return p.BeginCall(ctx, serviceName, operationName, callOptions)
}

// serve runs the listener to accept and manage new incoming connections, blocking
// until the channel is closed.
func (ch *Channel) serve() {
	acceptBackoff := 0 * time.Millisecond

	for {
		netConn, err := ch.mutable.l.Accept()
		if err != nil {
			// Backoff from new accepts if this is a temporary error
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if acceptBackoff == 0 {
					acceptBackoff = 5 * time.Millisecond
				} else {
					acceptBackoff *= 2
				}
				if max := 1 * time.Second; acceptBackoff > max {
					acceptBackoff = max
				}
				ch.log.Warnf("accept error: %v; retrying in %v", err, acceptBackoff)
				time.Sleep(acceptBackoff)
				continue
			} else {
				// Only log an error if this didn't happen due to a Close.
				if ch.State() >= ChannelStartClose {
					return
				}
				ch.log.Fatalf("unrecoverable accept error: %v; closing server", err)
				return
			}
		}

		acceptBackoff = 0

		// Register the connection in the peer once the channel is set up.
		events := connectionEvents{
			OnActive:           ch.incomingConnectionActive,
			OnCloseStateChange: ch.connectionCloseStateChange,
		}
		if _, err := ch.newInboundConnection(netConn, events, &ch.connectionOptions); err != nil {
			// Server is getting overloaded - begin rejecting new connections
			ch.log.Errorf("could not create new TChannelConnection for incoming conn: %v", err)
			netConn.Close()
			continue
		}
	}
}

// Ping sends a ping message to the given hostPort and waits for a response.
func (ch *Channel) Ping(ctx context.Context, hostPort string) error {
	peer := ch.Peers().GetOrAdd(hostPort)
	conn, err := peer.GetConnection(ctx)
	if err != nil {
		return err
	}

	return conn.ping(ctx)
}

// Logger returns the logger for this channel.
func (ch *Channel) Logger() Logger {
	return ch.log
}

// ServiceName returns the serviceName that this channel was created for.
func (ch *Channel) ServiceName() string {
	return ch.PeerInfo().ServiceName
}

// Connect connects the channel.
func (ch *Channel) Connect(ctx context.Context, hostPort string, connectionOptions *ConnectionOptions) (*Connection, error) {
	switch state := ch.State(); state {
	case ChannelClient, ChannelListening:
		break
	case ChannelStartClose:
		// We still allow outgoing connections during Close, but the connection has to immediately
		// be Closed after opening
	default:
		ch.log.Debugf("Connect rejecting new connection as state is %v", state)
		return nil, errInvalidStateForOp
	}

	events := connectionEvents{OnCloseStateChange: ch.connectionCloseStateChange}
	c, err := ch.newOutboundConnection(hostPort, events, connectionOptions)
	if err != nil {
		return nil, err
	}

	if err := c.sendInit(ctx); err != nil {
		return nil, err
	}

	ch.mutable.mut.Lock()
	ch.mutable.conns = append(ch.mutable.conns, c)
	chState := ch.mutable.state
	ch.mutable.mut.Unlock()

	// Any connections added after the channel is in StartClose should also be set to start close.
	if chState == ChannelStartClose {
		// TODO(prashant): If Connect is called, but no outgoing calls are made, then this connection
		// will block Close, as it will never get cleaned up.
		c.withStateLock(func() error {
			c.state = connectionStartClose
			return nil
		})
		c.log.Debugf("Channel is in start close, set connection to start close")
	}

	return c, err
}

// incomingConnectionActive adds a new active connection to our peer list.
func (ch *Channel) incomingConnectionActive(c *Connection) {
	c.log.Debugf("Add connection as an active peer for %v", c.remotePeerInfo.HostPort)
	p := ch.peers.GetOrAdd(c.remotePeerInfo.HostPort)
	p.AddConnection(c)

	ch.mutable.mut.Lock()
	ch.mutable.conns = append(ch.mutable.conns, c)
	ch.mutable.mut.Unlock()
}

// connectionCloseStateChange is called when a connection's close state changes.
func (ch *Channel) connectionCloseStateChange(c *Connection) {
	switch chState := ch.State(); chState {
	case ChannelStartClose, ChannelInboundClosed:
		ch.mutable.mut.RLock()
		minState := connectionClosed
		for _, c := range ch.mutable.conns {
			if s := c.readState(); s < minState {
				minState = s
			}
		}
		ch.mutable.mut.RUnlock()

		var updateTo ChannelState
		if minState >= connectionClosed {
			updateTo = ChannelClosed
		} else if minState >= connectionInboundClosed && chState == ChannelStartClose {
			updateTo = ChannelInboundClosed
		}

		if updateTo > 0 {
			ch.mutable.mut.Lock()
			ch.mutable.state = updateTo
			ch.mutable.mut.Unlock()
			chState = updateTo
		}

		c.log.Debugf("ConnectionCloseStateChange channel state = %v connection minState = %v",
			chState, minState)
	}
}

// Closed returns whether this channel has been closed with .Close()
func (ch *Channel) Closed() bool {
	return ch.State() == ChannelClosed
}

// State returns the current channel state.
func (ch *Channel) State() ChannelState {
	ch.mutable.mut.RLock()
	state := ch.mutable.state
	ch.mutable.mut.RUnlock()

	return state
}

// Close starts a graceful Close for the channel. This does not happen immediately:
// 1. This call closes the Listener and starts closing connections.
// 2. When all incoming connections are drainged, the connection blocks new outgoing calls.
// 3. When all connections are drainged, the channel's state is updated to Closed.
func (ch *Channel) Close() {
	ch.mutable.mut.Lock()

	if ch.mutable.l != nil {
		ch.mutable.l.Close()
	}
	ch.mutable.state = ChannelStartClose
	ch.mutable.mut.Unlock()

	ch.peers.Close()
}
<commit_msg>Add Peers method within the Registrar interface<commit_after>package tchannel

// Copyright (c) 2015 Uber Technologies, Inc.

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

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

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

var (
	errAlreadyListening  = errors.New("channel already listening")
	errInvalidStateForOp = errors.New("channel is in an invalid state for that operation")
)

const (
	ephemeralHostPort = "0.0.0.0:0"
)

// ChannelOptions are used to control parameters on a create a TChannel
type ChannelOptions struct {
	// Default Connection options
	DefaultConnectionOptions ConnectionOptions

	// The name of the process, for logging and reporting to peers
	ProcessName string

	// The logger to use for this channel
	Logger Logger

	// The reporter to use for reporting stats for this channel.
	StatsReporter StatsReporter

	// Trace reporter to use for this channel.
	TraceReporter TraceReporter
}

// ChannelState is the state of a channel.
type ChannelState int

const (
	// ChannelClient is a channel that can be used as a client.
	ChannelClient ChannelState = iota + 1

	// ChannelListening is a channel that is listening for new connnections.
	ChannelListening

	// ChannelStartClose is a channel that has received a Close request.
	// The channel is no longer listening, and all new incoming connections are rejected.
	ChannelStartClose

	// ChannelInboundClosed is a channel that has drained all incoming connections, but may
	// have outgoing connections. All incoming calls and new outgoing calls are rejected.
	ChannelInboundClosed

	// ChannelClosed is a channel that has closed completely.
	ChannelClosed
)

//go:generate stringer -type=ChannelState

// A Channel is a bi-directional connection to the peering and routing network.
// Applications can use a Channel to make service calls to remote peers via
// BeginCall, or to listen for incoming calls from peers.  Applications that
// want to receive requests should call one of Serve or ListenAndServe
// TODO(prashant): Shutdown all subchannels + peers when channel is closed.
type Channel struct {
	log               Logger
	commonStatsTags   map[string]string
	statsReporter     StatsReporter
	traceReporter     TraceReporter
	connectionOptions ConnectionOptions
	handlers          *handlerMap
	peers             *PeerList
	subChannels       *subChannelMap

	// mutable contains all the members of Channel which are mutable.
	mutable struct {
		mut      sync.RWMutex // protects members of the mutable struct.
		state    ChannelState
		peerInfo LocalPeerInfo // May be ephemeral if this is a client only channel
		l        net.Listener  // May be nil if this is a client only channel
		conns    []*Connection
	}
}

// NewChannel creates a new Channel.  The new channel can be used to send outbound requests
// to peers, but will not listen or handling incoming requests until one of ListenAndServe
// or Serve is called. The local service name should be passed to serviceName.
func NewChannel(serviceName string, opts *ChannelOptions) (*Channel, error) {
	if opts == nil {
		opts = &ChannelOptions{}
	}

	logger := opts.Logger
	if logger == nil {
		logger = NullLogger
	}

	traceReporter := opts.TraceReporter
	if traceReporter == nil {
		traceReporter = NullReporter
	}

	processName := opts.ProcessName
	if processName == "" {
		processName = fmt.Sprintf("%s[%d]", filepath.Base(os.Args[0]), os.Getpid())
	}

	statsReporter := opts.StatsReporter
	if statsReporter == nil {
		statsReporter = NullStatsReporter
	}

	ch := &Channel{
		connectionOptions: opts.DefaultConnectionOptions,
		log:               logger.WithFields(LogField{"service", serviceName}),
		statsReporter:     statsReporter,
		traceReporter:     traceReporter,
		handlers:          &handlerMap{},
		subChannels:       &subChannelMap{},
	}
	ch.mutable.peerInfo = LocalPeerInfo{
		PeerInfo: PeerInfo{
			ProcessName: processName,
			HostPort:    ephemeralHostPort,
		},
		ServiceName: serviceName,
	}
	ch.mutable.state = ChannelClient
	ch.peers = newPeerList(ch)
	ch.createCommonStats()
	return ch, nil
}

// Serve serves incoming requests using the provided listener.
// The local peer info is set synchronously, but the actual socket listening is done in
// a separate goroutine.
func (ch *Channel) Serve(l net.Listener) error {
	mutable := &ch.mutable
	mutable.mut.Lock()
	defer mutable.mut.Unlock()

	if mutable.l != nil {
		return errAlreadyListening
	}
	mutable.l = l

	if mutable.state != ChannelClient {
		return errInvalidStateForOp
	}
	mutable.state = ChannelListening

	mutable.peerInfo.HostPort = l.Addr().String()
	peerInfo := mutable.peerInfo
	ch.log.Debugf("%v (%v) listening on %v", peerInfo.ProcessName, peerInfo.ServiceName, peerInfo.HostPort)
	go ch.serve()
	return nil
}

// ListenAndServe listens on the given address and serves incoming requests.
// The port may be 0, in which case the channel will use an OS assigned port
// This method does not block as the handling of connections is done in a goroutine.
func (ch *Channel) ListenAndServe(hostPort string) error {
	mutable := &ch.mutable
	mutable.mut.RLock()

	if mutable.l != nil {
		mutable.mut.RUnlock()
		return errAlreadyListening
	}

	l, err := net.Listen("tcp", hostPort)
	if err != nil {
		mutable.mut.RUnlock()
		return err
	}

	mutable.mut.RUnlock()
	return ch.Serve(l)
}

// Registrar is the base interface for registering handlers on either the base
// Channel or the SubChannel
type Registrar interface {
	// ServiceName returns the service name that this Registrar is for.
	ServiceName() string

	// Register registers a handler for ServiceName and the given operation.
	Register(h Handler, operationName string)

	// Logger returns the logger for this Registrar.
	Logger() Logger

	// Peers returns the peerlist for this Registrar
	Peers() *PeerList
}

// Register registers a handler for a service+operation pair
func (ch *Channel) Register(h Handler, operationName string) {
	ch.handlers.register(h, ch.PeerInfo().ServiceName, operationName)
}

// PeerInfo returns the current peer info for the channel
func (ch *Channel) PeerInfo() LocalPeerInfo {
	ch.mutable.mut.RLock()
	defer ch.mutable.mut.RUnlock()

	return ch.mutable.peerInfo
}

func (ch *Channel) createCommonStats() {
	ch.commonStatsTags = map[string]string{
		"app":     ch.mutable.peerInfo.ProcessName,
		"service": ch.mutable.peerInfo.ServiceName,
	}
	host, err := os.Hostname()
	if err != nil {
		ch.log.Infof("channel failed to get host: %v", err)
		return
	}
	ch.commonStatsTags["host"] = host
	// TODO(prashant): Allow user to pass extra tags (such as cluster, version).
}

// GetSubChannel returns a SubChannel for the given service name. If the subchannel does not
// exist, it is created.
func (ch *Channel) GetSubChannel(serviceName string) *SubChannel {
	return ch.subChannels.getOrAdd(serviceName, ch)
}

// Peers returns the PeerList for the channel.
func (ch *Channel) Peers() *PeerList {
	return ch.peers
}

// BeginCall starts a new call to a remote peer, returning an OutboundCall that can
// be used to write the arguments of the call.
func (ch *Channel) BeginCall(ctx context.Context, hostPort, serviceName, operationName string, callOptions *CallOptions) (*OutboundCall, error) {
	p := ch.peers.GetOrAdd(hostPort)
	return p.BeginCall(ctx, serviceName, operationName, callOptions)
}

// serve runs the listener to accept and manage new incoming connections, blocking
// until the channel is closed.
func (ch *Channel) serve() {
	acceptBackoff := 0 * time.Millisecond

	for {
		netConn, err := ch.mutable.l.Accept()
		if err != nil {
			// Backoff from new accepts if this is a temporary error
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if acceptBackoff == 0 {
					acceptBackoff = 5 * time.Millisecond
				} else {
					acceptBackoff *= 2
				}
				if max := 1 * time.Second; acceptBackoff > max {
					acceptBackoff = max
				}
				ch.log.Warnf("accept error: %v; retrying in %v", err, acceptBackoff)
				time.Sleep(acceptBackoff)
				continue
			} else {
				// Only log an error if this didn't happen due to a Close.
				if ch.State() >= ChannelStartClose {
					return
				}
				ch.log.Fatalf("unrecoverable accept error: %v; closing server", err)
				return
			}
		}

		acceptBackoff = 0

		// Register the connection in the peer once the channel is set up.
		events := connectionEvents{
			OnActive:           ch.incomingConnectionActive,
			OnCloseStateChange: ch.connectionCloseStateChange,
		}
		if _, err := ch.newInboundConnection(netConn, events, &ch.connectionOptions); err != nil {
			// Server is getting overloaded - begin rejecting new connections
			ch.log.Errorf("could not create new TChannelConnection for incoming conn: %v", err)
			netConn.Close()
			continue
		}
	}
}

// Ping sends a ping message to the given hostPort and waits for a response.
func (ch *Channel) Ping(ctx context.Context, hostPort string) error {
	peer := ch.Peers().GetOrAdd(hostPort)
	conn, err := peer.GetConnection(ctx)
	if err != nil {
		return err
	}

	return conn.ping(ctx)
}

// Logger returns the logger for this channel.
func (ch *Channel) Logger() Logger {
	return ch.log
}

// ServiceName returns the serviceName that this channel was created for.
func (ch *Channel) ServiceName() string {
	return ch.PeerInfo().ServiceName
}

// Connect connects the channel.
func (ch *Channel) Connect(ctx context.Context, hostPort string, connectionOptions *ConnectionOptions) (*Connection, error) {
	switch state := ch.State(); state {
	case ChannelClient, ChannelListening:
		break
	case ChannelStartClose:
		// We still allow outgoing connections during Close, but the connection has to immediately
		// be Closed after opening
	default:
		ch.log.Debugf("Connect rejecting new connection as state is %v", state)
		return nil, errInvalidStateForOp
	}

	events := connectionEvents{OnCloseStateChange: ch.connectionCloseStateChange}
	c, err := ch.newOutboundConnection(hostPort, events, connectionOptions)
	if err != nil {
		return nil, err
	}

	if err := c.sendInit(ctx); err != nil {
		return nil, err
	}

	ch.mutable.mut.Lock()
	ch.mutable.conns = append(ch.mutable.conns, c)
	chState := ch.mutable.state
	ch.mutable.mut.Unlock()

	// Any connections added after the channel is in StartClose should also be set to start close.
	if chState == ChannelStartClose {
		// TODO(prashant): If Connect is called, but no outgoing calls are made, then this connection
		// will block Close, as it will never get cleaned up.
		c.withStateLock(func() error {
			c.state = connectionStartClose
			return nil
		})
		c.log.Debugf("Channel is in start close, set connection to start close")
	}

	return c, err
}

// incomingConnectionActive adds a new active connection to our peer list.
func (ch *Channel) incomingConnectionActive(c *Connection) {
	c.log.Debugf("Add connection as an active peer for %v", c.remotePeerInfo.HostPort)
	p := ch.peers.GetOrAdd(c.remotePeerInfo.HostPort)
	p.AddConnection(c)

	ch.mutable.mut.Lock()
	ch.mutable.conns = append(ch.mutable.conns, c)
	ch.mutable.mut.Unlock()
}

// connectionCloseStateChange is called when a connection's close state changes.
func (ch *Channel) connectionCloseStateChange(c *Connection) {
	switch chState := ch.State(); chState {
	case ChannelStartClose, ChannelInboundClosed:
		ch.mutable.mut.RLock()
		minState := connectionClosed
		for _, c := range ch.mutable.conns {
			if s := c.readState(); s < minState {
				minState = s
			}
		}
		ch.mutable.mut.RUnlock()

		var updateTo ChannelState
		if minState >= connectionClosed {
			updateTo = ChannelClosed
		} else if minState >= connectionInboundClosed && chState == ChannelStartClose {
			updateTo = ChannelInboundClosed
		}

		if updateTo > 0 {
			ch.mutable.mut.Lock()
			ch.mutable.state = updateTo
			ch.mutable.mut.Unlock()
			chState = updateTo
		}

		c.log.Debugf("ConnectionCloseStateChange channel state = %v connection minState = %v",
			chState, minState)
	}
}

// Closed returns whether this channel has been closed with .Close()
func (ch *Channel) Closed() bool {
	return ch.State() == ChannelClosed
}

// State returns the current channel state.
func (ch *Channel) State() ChannelState {
	ch.mutable.mut.RLock()
	state := ch.mutable.state
	ch.mutable.mut.RUnlock()

	return state
}

// Close starts a graceful Close for the channel. This does not happen immediately:
// 1. This call closes the Listener and starts closing connections.
// 2. When all incoming connections are drainged, the connection blocks new outgoing calls.
// 3. When all connections are drainged, the channel's state is updated to Closed.
func (ch *Channel) Close() {
	ch.mutable.mut.Lock()

	if ch.mutable.l != nil {
		ch.mutable.l.Close()
	}
	ch.mutable.state = ChannelStartClose
	ch.mutable.mut.Unlock()

	ch.peers.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"runtime"
	"strings"
	"time"

	"cloud.google.com/go/logging"
	"github.com/cloudfoundry-community/firehose-to-syslog/caching"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/config"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/filter"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/firehose"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/heartbeat"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/nozzle"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/stackdriver"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/version"
	"github.com/cloudfoundry-community/go-cfclient"
	"github.com/cloudfoundry/lager"
)

func main() {
	a := newApp()

	defer handleFatalError(a)

	producer := a.newProducer()
	consumer := a.newConsumer()

	errs, fhErrs := consumer.Start(producer)
	defer consumer.Stop()

	go func() {
		for err := range errs {
			a.logger.Error("nozzle", err)
		}
	}()

	fatalErr := <-fhErrs
	if fatalErr != nil {
		a.logger.Fatal("firehose", fatalErr)
	}
}

func handleFatalError(a *app) {
	if !a.c.DebugNozzle {
		return
	}

	if e := recover(); e != nil {
		stack := make([]byte, 1<<16)
		stackSize := runtime.Stack(stack, true)
		stackTrace := string(stack[:stackSize])

		payload := map[string]interface{}{
			"serviceContext": map[string]interface{}{
				"service": version.Name,
				"version": version.Release,
			},
			"message": stackTrace,
		}

		log := &stackdriver.Log{
			Payload:  payload,
			Labels:   map[string]string{},
			Severity: logging.Error,
		}

		logAdapter, err := a.newLogAdapter()
		if err == nil {
			logAdapter.PostLog(log)
			logAdapter.Flush()
		} else {
			a.logger.Error("error getting logAdapter", lager.Data{"err": err})
		}

		// Re-throw the error, we want to ensure it's logged directly to
		// stackdriver but we are not in a recoverable state.
		panic(e)
	}
}

func newApp() *app {
	logger := lager.NewLogger("stackdriver-nozzle")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.DEBUG))

	c, err := config.NewConfig()
	if err != nil {
		logger.Fatal("config", err)
	}

	logger.Info("arguments", c.ToData())

	trigger := time.NewTicker(time.Duration(c.HeartbeatRate) * time.Second).C
	heartbeater := heartbeat.NewHeartbeater(logger, trigger)

	cfConfig := &cfclient.Config{
		ApiAddress:        c.APIEndpoint,
		Username:          c.Username,
		Password:          c.Password,
		SkipSslValidation: c.SkipSSL}
	cfClient := cfclient.NewClient(cfConfig)

	var cachingClient caching.Caching
	if c.ResolveAppMetadata {
		cachingClient = caching.NewCachingBolt(cfClient, c.BoltDBPath)
	} else {
		cachingClient = caching.NewCachingEmpty()
	}
	cachingClient.CreateBucket()
	labelMaker := nozzle.NewLabelMaker(cachingClient)

	return &app{
		logger:      logger,
		c:           c,
		cfConfig:    cfConfig,
		cfClient:    cfClient,
		labelMaker:  labelMaker,
		heartbeater: heartbeater,
	}
}

type app struct {
	logger      lager.Logger
	c           *config.Config
	cfConfig    *cfclient.Config
	cfClient    *cfclient.Client
	labelMaker  nozzle.LabelMaker
	heartbeater heartbeat.Heartbeater
}

func (a *app) newProducer() firehose.Client {
	fhClient := firehose.NewClient(a.cfConfig, a.cfClient, a.c.SubscriptionID)

	producer, err := filter.New(fhClient, strings.Split(a.c.Events, ","), a.heartbeater)
	if err != nil {
		a.logger.Fatal("filter", err)
	}

	return producer
}

func (a *app) newConsumer() *nozzle.Nozzle {
	return &nozzle.Nozzle{
		LogSink:     a.newLogSink(),
		MetricSink:  a.newMetricSink(),
		Heartbeater: a.heartbeater,
	}
}

func (a *app) newLogSink() nozzle.Sink {
	logAdapter, logErrs := a.newLogAdapter()
	go func() {
		err := <-logErrs
		a.logger.Error("logAdapter", err)
	}()

	return nozzle.NewLogSink(a.labelMaker, logAdapter)
}

func (a *app) newLogAdapter() (stackdriver.LogAdapter, <-chan error) {
	return stackdriver.NewLogAdapter(
		a.c.ProjectID,
		a.c.BatchCount,
		time.Duration(a.c.BatchDuration)*time.Second,
		a.heartbeater,
	)
}

func (a *app) newMetricSink() nozzle.Sink {
	metricClient, err := stackdriver.NewMetricClient()
	if err != nil {
		a.logger.Fatal("metricClient", err)
	}

	metricAdapter, err := stackdriver.NewMetricAdapter(a.c.ProjectID, metricClient, a.heartbeater)
	if err != nil {
		a.logger.Error("metricAdapter", err)
	}

	metricBuffer, errs := stackdriver.NewMetricsBuffer(a.c.BatchCount, metricAdapter)
	go func() {
		for err = range errs {
			a.logger.Error("metricsBuffer", err)
		}
	}()

	return nozzle.NewMetricSink(a.labelMaker, metricBuffer, nozzle.NewUnitParser())
}
<commit_msg>Ignore the error channel<commit_after>package main

import (
	"os"
	"runtime"
	"strings"
	"time"

	"cloud.google.com/go/logging"
	"github.com/cloudfoundry-community/firehose-to-syslog/caching"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/config"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/filter"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/firehose"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/heartbeat"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/nozzle"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/stackdriver"
	"github.com/cloudfoundry-community/gcp-tools-release/src/stackdriver-nozzle/version"
	"github.com/cloudfoundry-community/go-cfclient"
	"github.com/cloudfoundry/lager"
)

func main() {
	a := newApp()

	defer handleFatalError(a)

	producer := a.newProducer()
	consumer := a.newConsumer()

	errs, fhErrs := consumer.Start(producer)
	defer consumer.Stop()

	go func() {
		for err := range errs {
			a.logger.Error("nozzle", err)
		}
	}()

	fatalErr := <-fhErrs
	if fatalErr != nil {
		a.logger.Fatal("firehose", fatalErr)
	}
}

func handleFatalError(a *app) {
	if !a.c.DebugNozzle {
		return
	}

	if e := recover(); e != nil {
		stack := make([]byte, 1<<16)
		stackSize := runtime.Stack(stack, true)
		stackTrace := string(stack[:stackSize])

		payload := map[string]interface{}{
			"serviceContext": map[string]interface{}{
				"service": version.Name,
				"version": version.Release,
			},
			"message": stackTrace,
		}

		log := &stackdriver.Log{
			Payload:  payload,
			Labels:   map[string]string{},
			Severity: logging.Error,
		}

		logAdapter, _ := a.newLogAdapter()
		logAdapter.PostLog(log)
		logAdapter.Flush()

		// Re-throw the error, we want to ensure it's logged directly to
		// stackdriver but we are not in a recoverable state.
		panic(e)
	}
}

func newApp() *app {
	logger := lager.NewLogger("stackdriver-nozzle")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.DEBUG))

	c, err := config.NewConfig()
	if err != nil {
		logger.Fatal("config", err)
	}

	logger.Info("arguments", c.ToData())

	trigger := time.NewTicker(time.Duration(c.HeartbeatRate) * time.Second).C
	heartbeater := heartbeat.NewHeartbeater(logger, trigger)

	cfConfig := &cfclient.Config{
		ApiAddress:        c.APIEndpoint,
		Username:          c.Username,
		Password:          c.Password,
		SkipSslValidation: c.SkipSSL}
	cfClient := cfclient.NewClient(cfConfig)

	var cachingClient caching.Caching
	if c.ResolveAppMetadata {
		cachingClient = caching.NewCachingBolt(cfClient, c.BoltDBPath)
	} else {
		cachingClient = caching.NewCachingEmpty()
	}
	cachingClient.CreateBucket()
	labelMaker := nozzle.NewLabelMaker(cachingClient)

	return &app{
		logger:      logger,
		c:           c,
		cfConfig:    cfConfig,
		cfClient:    cfClient,
		labelMaker:  labelMaker,
		heartbeater: heartbeater,
	}
}

type app struct {
	logger      lager.Logger
	c           *config.Config
	cfConfig    *cfclient.Config
	cfClient    *cfclient.Client
	labelMaker  nozzle.LabelMaker
	heartbeater heartbeat.Heartbeater
}

func (a *app) newProducer() firehose.Client {
	fhClient := firehose.NewClient(a.cfConfig, a.cfClient, a.c.SubscriptionID)

	producer, err := filter.New(fhClient, strings.Split(a.c.Events, ","), a.heartbeater)
	if err != nil {
		a.logger.Fatal("filter", err)
	}

	return producer
}

func (a *app) newConsumer() *nozzle.Nozzle {
	return &nozzle.Nozzle{
		LogSink:     a.newLogSink(),
		MetricSink:  a.newMetricSink(),
		Heartbeater: a.heartbeater,
	}
}

func (a *app) newLogSink() nozzle.Sink {
	logAdapter, logErrs := a.newLogAdapter()
	go func() {
		err := <-logErrs
		a.logger.Error("logAdapter", err)
	}()

	return nozzle.NewLogSink(a.labelMaker, logAdapter)
}

func (a *app) newLogAdapter() (stackdriver.LogAdapter, <-chan error) {
	return stackdriver.NewLogAdapter(
		a.c.ProjectID,
		a.c.BatchCount,
		time.Duration(a.c.BatchDuration)*time.Second,
		a.heartbeater,
	)
}

func (a *app) newMetricSink() nozzle.Sink {
	metricClient, err := stackdriver.NewMetricClient()
	if err != nil {
		a.logger.Fatal("metricClient", err)
	}

	metricAdapter, err := stackdriver.NewMetricAdapter(a.c.ProjectID, metricClient, a.heartbeater)
	if err != nil {
		a.logger.Error("metricAdapter", err)
	}

	metricBuffer, errs := stackdriver.NewMetricsBuffer(a.c.BatchCount, metricAdapter)
	go func() {
		for err = range errs {
			a.logger.Error("metricsBuffer", err)
		}
	}()

	return nozzle.NewMetricSink(a.labelMaker, metricBuffer, nozzle.NewUnitParser())
}
<|endoftext|>


<commit_before>package mux

//go:generate go run $GOPATH/src/v2ray.com/core/tools/generrorgen/main.go -pkg mux -path App,Proxyman,Mux

import (
	"context"
	"io"
	"sync"
	"time"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/ray"
)

const (
	maxTotal = 128
)

type ClientManager struct {
	access  sync.Mutex
	clients []*Client
	proxy   proxy.Outbound
	dialer  proxy.Dialer
	config  *proxyman.MultiplexingConfig
}

func NewClientManager(p proxy.Outbound, d proxy.Dialer, c *proxyman.MultiplexingConfig) *ClientManager {
	return &ClientManager{
		proxy:  p,
		dialer: d,
		config: c,
	}
}

func (m *ClientManager) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) error {
	m.access.Lock()
	defer m.access.Unlock()

	for _, client := range m.clients {
		if client.Dispatch(ctx, outboundRay) {
			return nil
		}
	}

	client, err := NewClient(m.proxy, m.dialer, m)
	if err != nil {
		return newError("failed to create client").Base(err)
	}
	m.clients = append(m.clients, client)
	client.Dispatch(ctx, outboundRay)
	return nil
}

func (m *ClientManager) onClientFinish() {
	m.access.Lock()
	defer m.access.Unlock()

	activeClients := make([]*Client, 0, len(m.clients))

	for _, client := range m.clients {
		if !client.Closed() {
			activeClients = append(activeClients, client)
		}
	}
	m.clients = activeClients
}

type Client struct {
	sessionManager *SessionManager
	inboundRay     ray.InboundRay
	ctx            context.Context
	cancel         context.CancelFunc
	manager        *ClientManager
	session2Remove chan uint16
	concurrency    uint32
}

var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527))

func NewClient(p proxy.Outbound, dialer proxy.Dialer, m *ClientManager) (*Client, error) {
	ctx, cancel := context.WithCancel(context.Background())
	ctx = proxy.ContextWithTarget(ctx, muxCoolDestination)
	pipe := ray.NewRay(ctx)
	go p.Process(ctx, pipe, dialer)
	c := &Client{
		sessionManager: NewSessionManager(),
		inboundRay:     pipe,
		ctx:            ctx,
		cancel:         cancel,
		manager:        m,
		session2Remove: make(chan uint16, 16),
		concurrency:    m.config.Concurrency,
	}
	go c.fetchOutput()
	go c.monitor()
	return c, nil
}

func (m *Client) Closed() bool {
	select {
	case <-m.ctx.Done():
		return true
	default:
		return false
	}
}

func (m *Client) monitor() {
	defer m.manager.onClientFinish()

	for {
		select {
		case <-m.ctx.Done():
			m.sessionManager.Close()
			m.inboundRay.InboundInput().Close()
			m.inboundRay.InboundOutput().CloseError()
			return
		case <-time.After(time.Second * 6):
			size := m.sessionManager.Size()
			if size == 0 && m.sessionManager.CloseIfNoSession() {
				m.cancel()
			}
		}
	}
}

func fetchInput(ctx context.Context, s *Session, output buf.Writer) {
	dest, _ := proxy.TargetFromContext(ctx)
	writer := &Writer{
		dest:   dest,
		id:     s.ID,
		writer: output,
	}
	defer writer.Close()
	defer s.CloseUplink()

	log.Trace(newError("dispatching request to ", dest))
	data, _ := s.input.ReadTimeout(time.Millisecond * 500)
	if err := writer.Write(data); err != nil {
		log.Trace(newError("failed to write first payload").Base(err))
		return
	}
	if err := buf.Copy(signal.BackgroundTimer(), s.input, writer); err != nil {
		log.Trace(newError("failed to fetch all input").Base(err))
	}
}

func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool {
	sm := m.sessionManager
	if sm.Size() >= int(m.concurrency) || sm.Count() >= maxTotal {
		return false
	}

	select {
	case <-m.ctx.Done():
		return false
	default:
	}

	s := sm.Allocate()
	if s == nil {
		return false
	}
	s.input = outboundRay.OutboundInput()
	s.output = outboundRay.OutboundOutput()
	go fetchInput(ctx, s, m.inboundRay.InboundInput())
	return true
}

func drain(reader *Reader) error {
	data, err := reader.Read()
	if err != nil {
		return err
	}
	data.Release()
	return nil
}

func pipe(reader *Reader, writer buf.Writer) error {
	data, err := reader.Read()
	if err != nil {
		return err
	}
	return writer.Write(data)
}

func (m *Client) handleStatueKeepAlive(meta *FrameMetadata, reader *Reader) error {
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (m *Client) handleStatusNew(meta *FrameMetadata, reader *Reader) error {
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (m *Client) handleStatusKeep(meta *FrameMetadata, reader *Reader) error {
	if !meta.Option.Has(OptionData) {
		return nil
	}

	if s, found := m.sessionManager.Get(meta.SessionID); found {
		return pipe(reader, s.output)
	}
	return drain(reader)
}

func (m *Client) handleStatusEnd(meta *FrameMetadata, reader *Reader) error {
	if s, found := m.sessionManager.Get(meta.SessionID); found {
		s.CloseDownlink()
		s.output.Close()
	}
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (m *Client) fetchOutput() {
	defer m.cancel()

	reader := NewReader(m.inboundRay.InboundOutput())
	for {
		meta, err := reader.ReadMetadata()
		if err != nil {
			if errors.Cause(err) != io.EOF {
				log.Trace(newError("failed to read metadata").Base(err))
			}
			break
		}

		switch meta.SessionStatus {
		case SessionStatusKeepAlive:
			err = m.handleStatueKeepAlive(meta, reader)
		case SessionStatusEnd:
			err = m.handleStatusEnd(meta, reader)
		case SessionStatusNew:
			err = m.handleStatusNew(meta, reader)
		case SessionStatusKeep:
			err = m.handleStatusKeep(meta, reader)
		default:
			log.Trace(newError("unknown status: ", meta.SessionStatus).AtWarning())
			return
		}

		if err != nil {
			log.Trace(newError("failed to process data").Base(err))
			return
		}
	}
}

type Server struct {
	dispatcher dispatcher.Interface
}

func NewServer(ctx context.Context) *Server {
	s := &Server{}
	space := app.SpaceFromContext(ctx)
	space.OnInitialize(func() error {
		d := dispatcher.FromSpace(space)
		if d == nil {
			return newError("no dispatcher in space")
		}
		s.dispatcher = d
		return nil
	})
	return s
}

func (s *Server) Dispatch(ctx context.Context, dest net.Destination) (ray.InboundRay, error) {
	if dest != muxCoolDestination {
		return s.dispatcher.Dispatch(ctx, dest)
	}

	ray := ray.NewRay(ctx)
	worker := &ServerWorker{
		dispatcher:     s.dispatcher,
		outboundRay:    ray,
		sessionManager: NewSessionManager(),
	}
	go worker.run(ctx)
	return ray, nil
}

type ServerWorker struct {
	dispatcher     dispatcher.Interface
	outboundRay    ray.OutboundRay
	sessionManager *SessionManager
}

func handle(ctx context.Context, s *Session, output buf.Writer) {
	writer := NewResponseWriter(s.ID, output)
	if err := buf.Copy(signal.BackgroundTimer(), s.input, writer); err != nil {
		log.Trace(newError("session ", s.ID, " ends: ").Base(err))
	}
	writer.Close()
	s.CloseDownlink()
}

func (w *ServerWorker) handleStatusKeepAlive(meta *FrameMetadata, reader *Reader) error {
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata, reader *Reader) error {
	log.Trace(newError("received request for ", meta.Target))
	inboundRay, err := w.dispatcher.Dispatch(ctx, meta.Target)
	if err != nil {
		if meta.Option.Has(OptionData) {
			drain(reader)
		}
		return newError("failed to dispatch request.").Base(err)
	}
	s := &Session{
		input:  inboundRay.InboundOutput(),
		output: inboundRay.InboundInput(),
		parent: w.sessionManager,
		ID:     meta.SessionID,
	}
	w.sessionManager.Add(s)
	go handle(ctx, s, w.outboundRay.OutboundOutput())
	if meta.Option.Has(OptionData) {
		return pipe(reader, s.output)
	}
	return nil
}

func (w *ServerWorker) handleStatusKeep(meta *FrameMetadata, reader *Reader) error {
	if !meta.Option.Has(OptionData) {
		return nil
	}
	if s, found := w.sessionManager.Get(meta.SessionID); found {
		return pipe(reader, s.output)
	}
	return drain(reader)
}

func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader *Reader) error {
	if s, found := w.sessionManager.Get(meta.SessionID); found {
		s.CloseUplink()
		s.output.Close()
	}
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (w *ServerWorker) run(ctx context.Context) {
	input := w.outboundRay.OutboundInput()
	reader := NewReader(input)

	defer w.sessionManager.Close()

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		meta, err := reader.ReadMetadata()
		if err != nil {
			log.Trace(newError("failed to read metadata").Base(err))
			return
		}

		switch meta.SessionStatus {
		case SessionStatusKeepAlive:
			err = w.handleStatusKeepAlive(meta, reader)
		case SessionStatusEnd:
			err = w.handleStatusEnd(meta, reader)
		case SessionStatusNew:
			err = w.handleStatusNew(ctx, meta, reader)
		case SessionStatusKeep:
			err = w.handleStatusKeep(meta, reader)
		default:
			log.Trace(newError("unknown status: ", meta.SessionStatus).AtWarning())
			return
		}

		if err != nil {
			log.Trace(newError("failed to process data").Base(err))
			return
		}
	}
}
<commit_msg>extend mux session timeout<commit_after>package mux

//go:generate go run $GOPATH/src/v2ray.com/core/tools/generrorgen/main.go -pkg mux -path App,Proxyman,Mux

import (
	"context"
	"io"
	"sync"
	"time"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/ray"
)

const (
	maxTotal = 128
)

type ClientManager struct {
	access  sync.Mutex
	clients []*Client
	proxy   proxy.Outbound
	dialer  proxy.Dialer
	config  *proxyman.MultiplexingConfig
}

func NewClientManager(p proxy.Outbound, d proxy.Dialer, c *proxyman.MultiplexingConfig) *ClientManager {
	return &ClientManager{
		proxy:  p,
		dialer: d,
		config: c,
	}
}

func (m *ClientManager) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) error {
	m.access.Lock()
	defer m.access.Unlock()

	for _, client := range m.clients {
		if client.Dispatch(ctx, outboundRay) {
			return nil
		}
	}

	client, err := NewClient(m.proxy, m.dialer, m)
	if err != nil {
		return newError("failed to create client").Base(err)
	}
	m.clients = append(m.clients, client)
	client.Dispatch(ctx, outboundRay)
	return nil
}

func (m *ClientManager) onClientFinish() {
	m.access.Lock()
	defer m.access.Unlock()

	activeClients := make([]*Client, 0, len(m.clients))

	for _, client := range m.clients {
		if !client.Closed() {
			activeClients = append(activeClients, client)
		}
	}
	m.clients = activeClients
}

type Client struct {
	sessionManager *SessionManager
	inboundRay     ray.InboundRay
	ctx            context.Context
	cancel         context.CancelFunc
	manager        *ClientManager
	session2Remove chan uint16
	concurrency    uint32
}

var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527))

func NewClient(p proxy.Outbound, dialer proxy.Dialer, m *ClientManager) (*Client, error) {
	ctx, cancel := context.WithCancel(context.Background())
	ctx = proxy.ContextWithTarget(ctx, muxCoolDestination)
	pipe := ray.NewRay(ctx)
	go p.Process(ctx, pipe, dialer)
	c := &Client{
		sessionManager: NewSessionManager(),
		inboundRay:     pipe,
		ctx:            ctx,
		cancel:         cancel,
		manager:        m,
		session2Remove: make(chan uint16, 16),
		concurrency:    m.config.Concurrency,
	}
	go c.fetchOutput()
	go c.monitor()
	return c, nil
}

func (m *Client) Closed() bool {
	select {
	case <-m.ctx.Done():
		return true
	default:
		return false
	}
}

func (m *Client) monitor() {
	defer m.manager.onClientFinish()

	for {
		select {
		case <-m.ctx.Done():
			m.sessionManager.Close()
			m.inboundRay.InboundInput().Close()
			m.inboundRay.InboundOutput().CloseError()
			return
		case <-time.After(time.Second * 16):
			size := m.sessionManager.Size()
			if size == 0 && m.sessionManager.CloseIfNoSession() {
				m.cancel()
			}
		}
	}
}

func fetchInput(ctx context.Context, s *Session, output buf.Writer) {
	dest, _ := proxy.TargetFromContext(ctx)
	writer := &Writer{
		dest:   dest,
		id:     s.ID,
		writer: output,
	}
	defer writer.Close()
	defer s.CloseUplink()

	log.Trace(newError("dispatching request to ", dest))
	data, _ := s.input.ReadTimeout(time.Millisecond * 500)
	if err := writer.Write(data); err != nil {
		log.Trace(newError("failed to write first payload").Base(err))
		return
	}
	if err := buf.Copy(signal.BackgroundTimer(), s.input, writer); err != nil {
		log.Trace(newError("failed to fetch all input").Base(err))
	}
}

func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool {
	sm := m.sessionManager
	if sm.Size() >= int(m.concurrency) || sm.Count() >= maxTotal {
		return false
	}

	select {
	case <-m.ctx.Done():
		return false
	default:
	}

	s := sm.Allocate()
	if s == nil {
		return false
	}
	s.input = outboundRay.OutboundInput()
	s.output = outboundRay.OutboundOutput()
	go fetchInput(ctx, s, m.inboundRay.InboundInput())
	return true
}

func drain(reader *Reader) error {
	data, err := reader.Read()
	if err != nil {
		return err
	}
	data.Release()
	return nil
}

func pipe(reader *Reader, writer buf.Writer) error {
	data, err := reader.Read()
	if err != nil {
		return err
	}
	return writer.Write(data)
}

func (m *Client) handleStatueKeepAlive(meta *FrameMetadata, reader *Reader) error {
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (m *Client) handleStatusNew(meta *FrameMetadata, reader *Reader) error {
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (m *Client) handleStatusKeep(meta *FrameMetadata, reader *Reader) error {
	if !meta.Option.Has(OptionData) {
		return nil
	}

	if s, found := m.sessionManager.Get(meta.SessionID); found {
		return pipe(reader, s.output)
	}
	return drain(reader)
}

func (m *Client) handleStatusEnd(meta *FrameMetadata, reader *Reader) error {
	if s, found := m.sessionManager.Get(meta.SessionID); found {
		s.CloseDownlink()
		s.output.Close()
	}
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (m *Client) fetchOutput() {
	defer m.cancel()

	reader := NewReader(m.inboundRay.InboundOutput())
	for {
		meta, err := reader.ReadMetadata()
		if err != nil {
			if errors.Cause(err) != io.EOF {
				log.Trace(newError("failed to read metadata").Base(err))
			}
			break
		}

		switch meta.SessionStatus {
		case SessionStatusKeepAlive:
			err = m.handleStatueKeepAlive(meta, reader)
		case SessionStatusEnd:
			err = m.handleStatusEnd(meta, reader)
		case SessionStatusNew:
			err = m.handleStatusNew(meta, reader)
		case SessionStatusKeep:
			err = m.handleStatusKeep(meta, reader)
		default:
			log.Trace(newError("unknown status: ", meta.SessionStatus).AtWarning())
			return
		}

		if err != nil {
			log.Trace(newError("failed to process data").Base(err))
			return
		}
	}
}

type Server struct {
	dispatcher dispatcher.Interface
}

func NewServer(ctx context.Context) *Server {
	s := &Server{}
	space := app.SpaceFromContext(ctx)
	space.OnInitialize(func() error {
		d := dispatcher.FromSpace(space)
		if d == nil {
			return newError("no dispatcher in space")
		}
		s.dispatcher = d
		return nil
	})
	return s
}

func (s *Server) Dispatch(ctx context.Context, dest net.Destination) (ray.InboundRay, error) {
	if dest != muxCoolDestination {
		return s.dispatcher.Dispatch(ctx, dest)
	}

	ray := ray.NewRay(ctx)
	worker := &ServerWorker{
		dispatcher:     s.dispatcher,
		outboundRay:    ray,
		sessionManager: NewSessionManager(),
	}
	go worker.run(ctx)
	return ray, nil
}

type ServerWorker struct {
	dispatcher     dispatcher.Interface
	outboundRay    ray.OutboundRay
	sessionManager *SessionManager
}

func handle(ctx context.Context, s *Session, output buf.Writer) {
	writer := NewResponseWriter(s.ID, output)
	if err := buf.Copy(signal.BackgroundTimer(), s.input, writer); err != nil {
		log.Trace(newError("session ", s.ID, " ends: ").Base(err))
	}
	writer.Close()
	s.CloseDownlink()
}

func (w *ServerWorker) handleStatusKeepAlive(meta *FrameMetadata, reader *Reader) error {
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata, reader *Reader) error {
	log.Trace(newError("received request for ", meta.Target))
	inboundRay, err := w.dispatcher.Dispatch(ctx, meta.Target)
	if err != nil {
		if meta.Option.Has(OptionData) {
			drain(reader)
		}
		return newError("failed to dispatch request.").Base(err)
	}
	s := &Session{
		input:  inboundRay.InboundOutput(),
		output: inboundRay.InboundInput(),
		parent: w.sessionManager,
		ID:     meta.SessionID,
	}
	w.sessionManager.Add(s)
	go handle(ctx, s, w.outboundRay.OutboundOutput())
	if meta.Option.Has(OptionData) {
		return pipe(reader, s.output)
	}
	return nil
}

func (w *ServerWorker) handleStatusKeep(meta *FrameMetadata, reader *Reader) error {
	if !meta.Option.Has(OptionData) {
		return nil
	}
	if s, found := w.sessionManager.Get(meta.SessionID); found {
		return pipe(reader, s.output)
	}
	return drain(reader)
}

func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader *Reader) error {
	if s, found := w.sessionManager.Get(meta.SessionID); found {
		s.CloseUplink()
		s.output.Close()
	}
	if meta.Option.Has(OptionData) {
		return drain(reader)
	}
	return nil
}

func (w *ServerWorker) run(ctx context.Context) {
	input := w.outboundRay.OutboundInput()
	reader := NewReader(input)

	defer w.sessionManager.Close()

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		meta, err := reader.ReadMetadata()
		if err != nil {
			log.Trace(newError("failed to read metadata").Base(err))
			return
		}

		switch meta.SessionStatus {
		case SessionStatusKeepAlive:
			err = w.handleStatusKeepAlive(meta, reader)
		case SessionStatusEnd:
			err = w.handleStatusEnd(meta, reader)
		case SessionStatusNew:
			err = w.handleStatusNew(ctx, meta, reader)
		case SessionStatusKeep:
			err = w.handleStatusKeep(meta, reader)
		default:
			log.Trace(newError("unknown status: ", meta.SessionStatus).AtWarning())
			return
		}

		if err != nil {
			log.Trace(newError("failed to process data").Base(err))
			return
		}
	}
}
<|endoftext|>


<commit_before>package goczmq

import (
	"fmt"
	"io"

	iomsg "github.com/taotetek/goczmq/iochunk/msg"
)

// ReadChunker accepts a socket and a chunkSize, and implements
// the ReaderFrom interface.
type ReadChunker struct {
	sock      *Sock
	chunkSize int64
}

// NewReadChunker takes a socket and a chunkSize and returns a new
// ReadChunker instance
func NewReadChunker(s *Sock, cs int64) *ReadChunker {
	return &ReadChunker{
		sock:      s,
		chunkSize: cs,
	}
}

// ReadFrom reads from an io.Reader into a []byte of chunkSize.
// It writes each chunk of data as a two frame message. The first
// frame is used to indicate if this is the last message or if
// there are more messages containing file data to come.
func (c *ReadChunker) ReadFrom(r io.Reader) (int64, error) {
	var total int64
	var n int
	var err error

	expect := []byte("MOAR")

	p := make([]byte, c.chunkSize)

	for err == nil {
		n, err = r.Read(p)

		msg := iomsg.NewChunk()
		msg.More = 1

		if err != nil {
			msg.More = 0
		}

		msg.Payload = p

		err = msg.Send(c.sock)
		if err != nil {
			return total, err
		}

		total += int64(n)
	}

	return total, err
}

// Destroy calls destroy on the underlying socket to clean it up
func (c *ReadChunker) Destroy() {
	c.sock.Destroy()
}

// WriteChunker accepts a socket and implements
// the WriterTo interface
type WriteChunker struct {
	sock *Sock
}

// NewWriteChunker takes a socket and returns a new
// WriteChunker instance
func NewWriteChunker(s *Sock) *WriteChunker {
	return &WriteChunker{
		sock: s,
	}
}

// WriteTo to reads each chunk message one at
// a time and writes  them to an io.Writer.
func (c *WriteChunker) WriteTo(w io.Writer) (int64, error) {
	var total int64
	var n int
	var err error
	var chunk [][]byte

	expect := []byte("MOAR")

	more := 1
	for more == 1 {
		transit, err := iomsg.Recv(c.sock)
		if err != nil {
			return total, fmt.Errorf("protocol error")
		}

		chunk := transit.(*iomsg.Chunk)
		n, err = w.Write(chunk.Payload)
		total += int64(n)
		if err != nil {
			return total, err
		}

		more = chunk.More
	}

	return total, err
}

// Destroy calls destroy on the underlying socket to clean it up
func (c *WriteChunker) Destroy() {
	c.sock.Destroy()
}
<commit_msg>reverting to naive implementation of chunker<commit_after>package goczmq

import (
	"bytes"
	"fmt"
	"io"
)

// ReadChunker accepts a socket and a chunkSize, and implements
// the ReaderFrom interface.
type ReadChunker struct {
	sock      *Sock
	chunkSize int64
}

// NewReadChunker takes a socket and a chunkSize and returns a new
// ReadChunker instance
func NewReadChunker(s *Sock, cs int64) *ReadChunker {
	return &ReadChunker{
		sock:      s,
		chunkSize: cs,
	}
}

// ReadFrom reads from an io.Reader into a []byte of chunkSize.
// It writes each chunk of data as a two frame message. The first
// frame is used to indicate if this is the last message or if
// there are more messages containing file data to come.
func (c *ReadChunker) ReadFrom(r io.Reader) (int64, error) {
	var total int64
	var n int
	var err error
	expect := []byte("MOAR")
	p := make([]byte, c.chunkSize)
	for err == nil {
		n, err = r.Read(p)
		if err != nil {
			expect = []byte("NOMOAR")
		}
		err := c.sock.SendFrame(expect, 1)
		if err != nil {
			return total, err
		}
		err = c.sock.SendFrame(p[:n], 0)
		if err != nil {
			return total, err
		}
		total += int64(n)
	}
	return total, err
}

// Destroy calls destroy on the underlying socket to clean it up
func (c *ReadChunker) Destroy() {
	c.sock.Destroy()
}

// WriteChunker accepts a socket and implements
// the WriterTo interface
type WriteChunker struct {
	sock *Sock
}

// NewWriteChunker takes a socket and returns a new
// WriteChunker instance
func NewWriteChunker(s *Sock) *WriteChunker {
	return &WriteChunker{
		sock: s,
	}
}

// WriteTo to reads each chunk message one at
// a time and writes them to an io.Writer.
func (c *WriteChunker) WriteTo(w io.Writer) (int64, error) {
	var total int64
	var n int
	var err error
	var chunk [][]byte
	expect := []byte("MOAR")
	for bytes.Compare(expect, []byte("MOAR")) == 0 {
		chunk, err = c.sock.RecvMessage()
		if len(chunk) != 2 {
			return total, fmt.Errorf("protocol error")
		}
		expect = chunk[0]
		if bytes.Compare(expect, []byte("MOAR")) != 0 &&
			bytes.Compare(expect, []byte("NOMOAR")) != 0 {
			return total, fmt.Errorf("protocol error")
		}
		n, err = w.Write(chunk[1])
		total += int64(n)
		if err != nil {
			return total, err
		}
	}
	return total, err
}

// Destroy calls destroy on the underlying socket to clean it up
func (c *WriteChunker) Destroy() {
	c.sock.Destroy()
}
<|endoftext|>


<commit_before>/**
 * This file is covered by the AGPLv3 license, which can be found at the LICENSE file in the root of this project.
 * @copyright 2013 The Gorilla WebSocket Authors
 * @copyright 2017 subtitulamos.tv
 */

package main

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// connection is an middleman between the websocket connection and the hub.
type connection struct {
	// The websocket connection.
	ws *websocket.Conn

	// The hub it belongs to
	hub *hub

	// Buffered channel of outbound messages.
	send chan []byte
}

// write writes a message with the given message type and payload.
func (c *connection) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

// readPump listens to messages from client (such as disconnection)
func (c *connection) readPump() {
	defer func() {
		c.hub.leave <- c
		c.ws.Close()
	}()
	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, _, err := c.ws.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("error: %v", err)
			}
			break
		}
	}
}

// writePump pumps messages from the hub to the websocket connection.
func (c *connection) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.write(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

// serveWs handles websocket requests from the peer.
func serveWs(w http.ResponseWriter, r *http.Request) {
	subID, _ := strconv.Atoi(r.URL.Query().Get("subID"))
	token := r.URL.Query().Get("token")

	log.Printf("Client logging in with token <%s> on sub#%d", token, subID)
	if subID == 0 || len(token) == 0 {
		log.Println("Invalid access arguments")
		return
	}

	v, err := redisClient.Get(fmt.Sprintf("authtok-%s-%s", redisEnvPrefix, token)).Result()
	if err != nil {
		log.Printf("(%s, %d) pair fail", token, subID)
		return
	}

	expectedSubID, err := strconv.Atoi(v)
	if err != nil || expectedSubID != subID {
		log.Printf("(%s, %d) pair fail", token, subID)
		return
	}

	log.Printf("(%s, %d) pair OK", token, subID)

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	sub, ok := subtitles[subID]
	if !ok {
		sub = newSubtitle(subID)
		subtitles[subID] = sub

		log.Printf("Created subtitle instance for %d", subID)
	}

	c := &connection{send: make(chan []byte, 256), ws: ws, hub: sub.hub}
	sub.hub.join <- c
	go c.writePump()
	c.readPump()
}
<commit_msg>Improve logging when keyv pair is rejected while setting up a WS<commit_after>/**
 * This file is covered by the AGPLv3 license, which can be found at the LICENSE file in the root of this project.
 * @copyright 2013 The Gorilla WebSocket Authors
 * @copyright 2017 subtitulamos.tv
 */

package main

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// connection is an middleman between the websocket connection and the hub.
type connection struct {
	// The websocket connection.
	ws *websocket.Conn

	// The hub it belongs to
	hub *hub

	// Buffered channel of outbound messages.
	send chan []byte
}

// write writes a message with the given message type and payload.
func (c *connection) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

// readPump listens to messages from client (such as disconnection)
func (c *connection) readPump() {
	defer func() {
		c.hub.leave <- c
		c.ws.Close()
	}()
	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, _, err := c.ws.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Printf("error: %v", err)
			}
			break
		}
	}
}

// writePump pumps messages from the hub to the websocket connection.
func (c *connection) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.write(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

// serveWs handles websocket requests from the peer.
func serveWs(w http.ResponseWriter, r *http.Request) {
	subID, _ := strconv.Atoi(r.URL.Query().Get("subID"))
	token := r.URL.Query().Get("token")

	log.Printf("Client logging in with token <%s> on sub#%d", token, subID)
	if subID == 0 || len(token) == 0 {
		log.Println("Invalid access arguments")
		return
	}

	auths := fmt.Sprintf("authtok-%s-%s", redisEnvPrefix, token)
	v, err := redisClient.Get(auths).Result()
	if err != nil {
		log.Printf("(%s, %d) pair fail, not found <%s>. err: %v", token, subID, auths, err)
		return
	}

	expectedSubID, err := strconv.Atoi(v)
	if err != nil || expectedSubID != subID {
		log.Printf("(%s, %d) pair fail, wrong sub ID", token, subID)
		return
	}

	log.Printf("(%s, %d) pair OK", token, subID)

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	sub, ok := subtitles[subID]
	if !ok {
		sub = newSubtitle(subID)
		subtitles[subID] = sub

		log.Printf("Created subtitle instance for %d", subID)
	}

	c := &connection{send: make(chan []byte, 256), ws: ws, hub: sub.hub}
	sub.hub.join <- c
	go c.writePump()
	c.readPump()
}
<|endoftext|>