<commit_before>package exportertools

import(
    "github.com/prometheus/common/log"
    "github.com/prometheus/client_golang/prometheus"
    "time"
)

/*
BaseExporter provides convinient caching and metrics processing.
BaseExporter should be embedded into custom exporter struct
*/
type BaseExporter struct {
    Control chan bool
    Name    string
    Cache   *Cache
    metrics []MetricCollector
    Labels  map[string]string
}

func NewBaseExporter(name string, ttl int, labels map[string]string) *BaseExporter {
    e := BaseExporter{
        Control:  make(chan bool, 1),
        Name:     name,
        Cache:    NewCache(ttl),
        metrics:  []MetricCollector{},
        Labels:   labels,
    }
    return &e
}

func (e *BaseExporter) AddCollector(m MetricCollector) {
    e.metrics = append(e.metrics, m)
}

// Process runs metric collection in async mode
func (e *BaseExporter) Process() {
    defer log.Debugf("Instance terminated. Bye-bye from %v", e.Name, e)

    var channels []chan bool

    for _, _ = range e.metrics {
        channels = append(channels, make(chan bool))
    }

    for id, mc := range e.metrics {
        log.Debugf("Started Process() for %v", e.Name)
        ticker := time.NewTicker(e.Cache.TTL)
        go func(done chan bool, m MetricCollector) {
            defer close(done)
            for {
                select {
                case <-ticker.C:
                    go func() {
                        log.Debugf("Call Collect() for %v", e.Name)
                        metrics, err := m.Collect()
                        if err != nil {
                            log.Errorf("Error occured during Collect() for %v. error: %v", e.Name, err)
                        }
                        for idx := range metrics {
                            e.Cache.Set(metrics[idx])
                        }
                    }()
                case <-done:
                    log.Debugf("Exit from child of Process() %v", e.Name)
                    return
                }
            }
        }(channels[id], mc)
    }
    msg := <- e.Control
    go broadcastMessage(msg, channels)
}

func (e *BaseExporter) Close() error {
    defer close(e.Control)

    e.Control<- true
    log.Debugf("Stop processing metric for %v", e.Labels)
    err := Unregister(e)
    return err
}

func (e *BaseExporter) Setup() error {
    log.Debugf("Default Setup() func is used. It does nothing")
    return nil
}

// Satisfies Exporter interface and calls GenericDescribe which works with cache
func (e *BaseExporter) Describe(ch chan<- *prometheus.Desc) {
    GenericDescribe(e, ch)
}

// Satisfies Exporter interface and calls GenericCollect which works with cache
func (e *BaseExporter) Collect(ch chan<- prometheus.Metric) {
    GenericCollect(e, ch)
}

func broadcastMessage(msg bool, channels []chan bool) {
    for _, ch := range channels {
        ch <- msg
    }
}
<commit_msg>change field name, public access<commit_after>package exportertools

import(
    "github.com/prometheus/common/log"
    "github.com/prometheus/client_golang/prometheus"
    "time"
)

/*
BaseExporter provides convinient caching and metrics processing.
BaseExporter should be embedded into custom exporter struct
*/
type BaseExporter struct {
    Control chan bool
    Name    string
    Cache   *Cache
    MetricCollectors []MetricCollector
    Labels  map[string]string
}

func NewBaseExporter(name string, ttl int, labels map[string]string) *BaseExporter {
    e := BaseExporter{
        Control:  make(chan bool, 1),
        Name:     name,
        Cache:    NewCache(ttl),
        MetricCollectors:  []MetricCollector{},
        Labels:   labels,
    }
    return &e
}

func (e *BaseExporter) AddCollector(m MetricCollector) {
    e.MetricCollectors = append(e.MetricCollectors, m)
}

// Process runs metric collection in async mode
func (e *BaseExporter) Process() {
    defer log.Debugf("Instance terminated. Bye-bye from %v", e.Name, e)

    var channels []chan bool

    for _, _ = range e.MetricCollectors {
        channels = append(channels, make(chan bool))
    }

    for id, mc := range e.MetricCollectors {
        log.Debugf("Started Process() for %v", e.Name)
        ticker := time.NewTicker(e.Cache.TTL)
        go func(done chan bool, m MetricCollector) {
            defer close(done)
            for {
                select {
                case <-ticker.C:
                    go func() {
                        log.Debugf("Call Collect() for %v", e.Name)
                        metrics, err := m.Collect()
                        if err != nil {
                            log.Errorf("Error occured during Collect() for %v. error: %v", e.Name, err)
                        }
                        for idx := range metrics {
                            e.Cache.Set(metrics[idx])
                        }
                    }()
                case <-done:
                    log.Debugf("Exit from child of Process() %v", e.Name)
                    return
                }
            }
        }(channels[id], mc)
    }
    msg := <- e.Control
    go broadcastMessage(msg, channels)
}

func (e *BaseExporter) Close() error {
    defer close(e.Control)

    e.Control<- true
    log.Debugf("Stop processing metric for %v", e.Labels)
    err := Unregister(e)
    return err
}

func (e *BaseExporter) Setup() error {
    log.Debugf("Default Setup() func is used. It does nothing")
    return nil
}

// Satisfies Exporter interface and calls GenericDescribe which works with cache
func (e *BaseExporter) Describe(ch chan<- *prometheus.Desc) {
    GenericDescribe(e, ch)
}

// Satisfies Exporter interface and calls GenericCollect which works with cache
func (e *BaseExporter) Collect(ch chan<- prometheus.Metric) {
    GenericCollect(e, ch)
}

func broadcastMessage(msg bool, channels []chan bool) {
    for _, ch := range channels {
        ch <- msg
    }
}
<|endoftext|>


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

package mtls

import (
	"fmt"

	"istio.io/api/authentication/v1alpha1"

	"istio.io/istio/galley/pkg/config/resource"

	"istio.io/istio/galley/pkg/config/analysis/analyzers/util"
)

// TargetService is a simple struct type for representing a service
// targeted by an Authentication policy.
type TargetService struct {
	fQDN       string
	portNumber uint32
}

// NewTargetServiceWithPortNumber creates a new TargetService using the specified
// fqdn and portNumber.
func NewTargetServiceWithPortNumber(fqdn string, portNumber uint32) TargetService {
	return TargetService{fQDN: fqdn, portNumber: portNumber}
}

// NewTargetService creates a new TargetService using the specified fqdn. Because no
// port is specified, this implicitly represents the service bound to any port.
func NewTargetService(fqdn string) TargetService {
	return TargetService{fQDN: fqdn}
}

// FQDN is the fully-qualified domain name for the service (e.g.
// foobar.my-namespace.svc.cluster.local).
func (w TargetService) FQDN() string {
	return w.fQDN
}

// PortNumber is the port used by the service.
func (w TargetService) PortNumber() uint32 {
	return w.portNumber
}

func (w TargetService) String() string {
	if w.PortNumber() != 0 {
		return fmt.Sprintf("%s:%d", w.fQDN, w.portNumber)
	}
	return w.FQDN()
}

// PolicyChecker allows callers to add a set of v1alpha1.Policy objects in the
// mesh. Once these are loaded, you can query whether or not a specific
// TargetService will require MTLS when an incoming connection occurs using the
// IsServiceMTLSEnforced() call.
type PolicyChecker struct {
	meshMTLSModeAndResource ModeAndResource

	namespaceToMTLSMode        map[string]ModeAndResource
	serviceToMTLSMode          map[TargetService]ModeAndResource
	fqdnToPortNameToPortNumber map[string]map[string]uint32
}

// Mode is a special type used to distinguish between MTLS being off
// (unsupported), permissive (supported but not required), and strict (required).
type Mode int

const (
	// ModePlaintext means MTLS is off (unsupported).
	ModePlaintext Mode = iota
	// ModePermissive means MTLS is permissive (supported but not required)
	ModePermissive
	// ModeStrict means MTLS is strict (required)
	ModeStrict
)

func (m Mode) String() string {
	switch m {
	case ModePlaintext:
		return "Plaintext"
	case ModePermissive:
		return "Permissive"
	case ModeStrict:
		return "Strict"
	default:
		return "UNKNOWN"
	}
}

// ModeAndResource is a simple tuple type of mode and the resource that
// specified the mode.
type ModeAndResource struct {
	MTLSMode Mode
	Resource *resource.Entry
}

// NewPolicyChecker creates a new PolicyChecker instance.
func NewPolicyChecker(fqdnToPortNameToPortNumber map[string]map[string]uint32) *PolicyChecker {
	return &PolicyChecker{
		namespaceToMTLSMode:        make(map[string]ModeAndResource),
		serviceToMTLSMode:          make(map[TargetService]ModeAndResource),
		fqdnToPortNameToPortNumber: fqdnToPortNameToPortNumber,
	}
}

// AddMeshPolicy adds a mesh-level policy to the checker. Note that there can
// only be at most one mesh level policy in effect.
func (pc *PolicyChecker) AddMeshPolicy(r *resource.Entry, p *v1alpha1.Policy) error {
	mode, err := parsePolicyMTLSMode(p)
	if err != nil {
		return err
	}
	pc.meshMTLSModeAndResource = ModeAndResource{Resource: r, MTLSMode: mode}
	return nil
}

// MeshPolicy returns the current recognized MeshPolicy (as added by AddMeshPolicy).
func (pc *PolicyChecker) MeshPolicy() ModeAndResource {
	return pc.meshMTLSModeAndResource
}

type NamedPortInPolicyNotFoundError struct {
	PortName     string
	PolicyOrigin string
	FQDN         string
}

func (e NamedPortInPolicyNotFoundError) Error() string {
	return fmt.Sprintf("named port '%s' not found for fqdn '%s', unable to analyze policy '%s'", e.PortName, e.FQDN, e.PolicyOrigin)
}

// AddPolicy adds a new policy object to the PolicyChecker to use when later
// determining if a service is MTLS-enforced. The namespace of the policy is
// also provided as some policies can target the local namespace.
//
// If the Policy uses a named port, and the port cannot be looked up in the map
// provided to NewPolicyChecker, then an error of type
// NamedPortInPolicyNotFoundError is returned.
func (pc *PolicyChecker) AddPolicy(r *resource.Entry, p *v1alpha1.Policy) error {
	mode, err := parsePolicyMTLSMode(p)
	if err != nil {
		return err
	}
	modeAndResource := ModeAndResource{Resource: r, MTLSMode: mode}
	namespace, _ := r.Metadata.Name.InterpretAsNamespaceAndName()
	if len(p.Targets) == 0 {
		// Rule targets the namespace.
		pc.namespaceToMTLSMode[namespace] = modeAndResource
		return nil
	}
	// Discover the targeted service and take note. Should normalize.
	for _, target := range p.Targets {
		fqdn := util.ConvertHostToFQDN(namespace, target.Name)

		if len(target.Ports) == 0 {
			// Policy targets all ports on service
			pc.serviceToMTLSMode[NewTargetService(fqdn)] = modeAndResource
		}

		for _, port := range target.Ports {
			if port.GetName() != "" {
				// Look up the port number for the name. If we can't find it, we
				// need to complain about a different error
				// TODO handle missing reference error.
				if _, ok := pc.fqdnToPortNameToPortNumber[fqdn]; !ok {
					return NamedPortInPolicyNotFoundError{
						PortName:     port.GetName(),
						PolicyOrigin: r.Origin.FriendlyName(),
						FQDN:         fqdn,
					}
				}
				portNumber := pc.fqdnToPortNameToPortNumber[fqdn][port.GetName()]
				if portNumber == 0 {
					return NamedPortInPolicyNotFoundError{
						PortName:     port.GetName(),
						PolicyOrigin: r.Origin.FriendlyName(),
						FQDN:         fqdn,
					}
				}

				pc.serviceToMTLSMode[NewTargetServiceWithPortNumber(fqdn, portNumber)] = modeAndResource
			} else if port.GetNumber() != 0 {
				pc.serviceToMTLSMode[NewTargetServiceWithPortNumber(fqdn, port.GetNumber())] = modeAndResource
			} else {
				// Unhandled case!
				return fmt.Errorf("policy has a port with no name/number for target %s", target.Name)
			}
		}
	}

	return nil
}

// IsServiceMTLSEnforced returns true if a service requires incoming
// connections to use MTLS, or false if MTLS is not a hard-requirement (e.g.
// mode is permissive, peerIsOptional is true, etc). Only call this after adding
// all policy resources in effect via AddPolicy or AddMeshPolicy.
func (pc *PolicyChecker) IsServiceMTLSEnforced(w TargetService) (ModeAndResource, error) {
	// TODO support understanding port name -> port number mappings
	var modeAndResource ModeAndResource
	modeAndResource = pc.serviceToMTLSMode[w]
	if modeAndResource.Resource != nil {
		return modeAndResource, nil
	}

	// Try checking if its enforced on any ports
	serviceNoPort := NewTargetService(w.FQDN())
	modeAndResource = pc.serviceToMTLSMode[serviceNoPort]
	if modeAndResource.Resource != nil {
		return modeAndResource, nil
	}

	// Check if enforced on namespace
	namespace, _ := util.GetResourceNameFromHost("", w.FQDN()).InterpretAsNamespaceAndName()
	if namespace == "" {
		return ModeAndResource{}, fmt.Errorf("unable to extract namespace from fqdn: %s", w.FQDN())
	}
	modeAndResource = pc.namespaceToMTLSMode[namespace]
	if modeAndResource.Resource != nil {
		return modeAndResource, nil
	}

	// Finally, defer to mesh level policy. No need to check for a nil resource
	// since the default value is correct.
	return pc.meshMTLSModeAndResource, nil
}

// TargetServices returns the list of services known to the policy checker (in
// no particular order).
func (pc *PolicyChecker) TargetServices() []TargetService {
	tss := make([]TargetService, len(pc.serviceToMTLSMode))
	i := 0
	for ts := range pc.serviceToMTLSMode {
		tss[i] = ts
		i++
	}

	return tss
}

// parsePolicyMTLSMode is a helper function to determine what mtls mode a Policy
// implies.
func parsePolicyMTLSMode(p *v1alpha1.Policy) (Mode, error) {
	for _, peer := range p.Peers {
		mtlsParams, ok := peer.Params.(*v1alpha1.PeerAuthenticationMethod_Mtls)
		if !ok {
			// Only looking for mtls methods
			continue
		}

		var mode Mode
		if mtlsParams.Mtls == nil {
			mode = ModeStrict
		} else {
			switch mtlsParams.Mtls.GetMode() {
			case v1alpha1.MutualTls_PERMISSIVE:
				mode = ModePermissive
			case v1alpha1.MutualTls_STRICT:
				mode = ModeStrict
			default:
				// Shouldn't happen!
				return mode, fmt.Errorf("unknown MTLS mode when analyzing policy: %s", mtlsParams.Mtls.GetMode().String())
			}
		}
		// Now check for modifiers that might downgrade strict to permissive.
		if mode == ModeStrict && mtlsParams.Mtls != nil && mtlsParams.Mtls.AllowTls {
			mode = ModePermissive
		}
		if mode == ModeStrict && p.PeerIsOptional {
			mode = ModePermissive
		}
		return mode, nil
	}

	// No MTLS configuration found
	return ModePlaintext, nil
}
<commit_msg>Use fqdn rather than fQDN for non-exported field (#19249)<commit_after>// Copyright 2019 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mtls

import (
	"fmt"

	"istio.io/api/authentication/v1alpha1"

	"istio.io/istio/galley/pkg/config/resource"

	"istio.io/istio/galley/pkg/config/analysis/analyzers/util"
)

// TargetService is a simple struct type for representing a service
// targeted by an Authentication policy.
type TargetService struct {
	fqdn       string
	portNumber uint32
}

// NewTargetServiceWithPortNumber creates a new TargetService using the specified
// fqdn and portNumber.
func NewTargetServiceWithPortNumber(fqdn string, portNumber uint32) TargetService {
	return TargetService{fqdn: fqdn, portNumber: portNumber}
}

// NewTargetService creates a new TargetService using the specified fqdn. Because no
// port is specified, this implicitly represents the service bound to any port.
func NewTargetService(fqdn string) TargetService {
	return TargetService{fqdn: fqdn}
}

// FQDN is the fully-qualified domain name for the service (e.g.
// foobar.my-namespace.svc.cluster.local).
func (w TargetService) FQDN() string {
	return w.fqdn
}

// PortNumber is the port used by the service.
func (w TargetService) PortNumber() uint32 {
	return w.portNumber
}

func (w TargetService) String() string {
	if w.PortNumber() != 0 {
		return fmt.Sprintf("%s:%d", w.fqdn, w.portNumber)
	}
	return w.fqdn
}

// PolicyChecker allows callers to add a set of v1alpha1.Policy objects in the
// mesh. Once these are loaded, you can query whether or not a specific
// TargetService will require MTLS when an incoming connection occurs using the
// IsServiceMTLSEnforced() call.
type PolicyChecker struct {
	meshMTLSModeAndResource ModeAndResource

	namespaceToMTLSMode        map[string]ModeAndResource
	serviceToMTLSMode          map[TargetService]ModeAndResource
	fqdnToPortNameToPortNumber map[string]map[string]uint32
}

// Mode is a special type used to distinguish between MTLS being off
// (unsupported), permissive (supported but not required), and strict (required).
type Mode int

const (
	// ModePlaintext means MTLS is off (unsupported).
	ModePlaintext Mode = iota
	// ModePermissive means MTLS is permissive (supported but not required)
	ModePermissive
	// ModeStrict means MTLS is strict (required)
	ModeStrict
)

func (m Mode) String() string {
	switch m {
	case ModePlaintext:
		return "Plaintext"
	case ModePermissive:
		return "Permissive"
	case ModeStrict:
		return "Strict"
	default:
		return "UNKNOWN"
	}
}

// ModeAndResource is a simple tuple type of mode and the resource that
// specified the mode.
type ModeAndResource struct {
	MTLSMode Mode
	Resource *resource.Entry
}

// NewPolicyChecker creates a new PolicyChecker instance.
func NewPolicyChecker(fqdnToPortNameToPortNumber map[string]map[string]uint32) *PolicyChecker {
	return &PolicyChecker{
		namespaceToMTLSMode:        make(map[string]ModeAndResource),
		serviceToMTLSMode:          make(map[TargetService]ModeAndResource),
		fqdnToPortNameToPortNumber: fqdnToPortNameToPortNumber,
	}
}

// AddMeshPolicy adds a mesh-level policy to the checker. Note that there can
// only be at most one mesh level policy in effect.
func (pc *PolicyChecker) AddMeshPolicy(r *resource.Entry, p *v1alpha1.Policy) error {
	mode, err := parsePolicyMTLSMode(p)
	if err != nil {
		return err
	}
	pc.meshMTLSModeAndResource = ModeAndResource{Resource: r, MTLSMode: mode}
	return nil
}

// MeshPolicy returns the current recognized MeshPolicy (as added by AddMeshPolicy).
func (pc *PolicyChecker) MeshPolicy() ModeAndResource {
	return pc.meshMTLSModeAndResource
}

type NamedPortInPolicyNotFoundError struct {
	PortName     string
	PolicyOrigin string
	FQDN         string
}

func (e NamedPortInPolicyNotFoundError) Error() string {
	return fmt.Sprintf("named port '%s' not found for fqdn '%s', unable to analyze policy '%s'", e.PortName, e.FQDN, e.PolicyOrigin)
}

// AddPolicy adds a new policy object to the PolicyChecker to use when later
// determining if a service is MTLS-enforced. The namespace of the policy is
// also provided as some policies can target the local namespace.
//
// If the Policy uses a named port, and the port cannot be looked up in the map
// provided to NewPolicyChecker, then an error of type
// NamedPortInPolicyNotFoundError is returned.
func (pc *PolicyChecker) AddPolicy(r *resource.Entry, p *v1alpha1.Policy) error {
	mode, err := parsePolicyMTLSMode(p)
	if err != nil {
		return err
	}
	modeAndResource := ModeAndResource{Resource: r, MTLSMode: mode}
	namespace, _ := r.Metadata.Name.InterpretAsNamespaceAndName()
	if len(p.Targets) == 0 {
		// Rule targets the namespace.
		pc.namespaceToMTLSMode[namespace] = modeAndResource
		return nil
	}
	// Discover the targeted service and take note. Should normalize.
	for _, target := range p.Targets {
		fqdn := util.ConvertHostToFQDN(namespace, target.Name)

		if len(target.Ports) == 0 {
			// Policy targets all ports on service
			pc.serviceToMTLSMode[NewTargetService(fqdn)] = modeAndResource
		}

		for _, port := range target.Ports {
			if port.GetName() != "" {
				// Look up the port number for the name. If we can't find it, we
				// need to complain about a different error
				// TODO handle missing reference error.
				if _, ok := pc.fqdnToPortNameToPortNumber[fqdn]; !ok {
					return NamedPortInPolicyNotFoundError{
						PortName:     port.GetName(),
						PolicyOrigin: r.Origin.FriendlyName(),
						FQDN:         fqdn,
					}
				}
				portNumber := pc.fqdnToPortNameToPortNumber[fqdn][port.GetName()]
				if portNumber == 0 {
					return NamedPortInPolicyNotFoundError{
						PortName:     port.GetName(),
						PolicyOrigin: r.Origin.FriendlyName(),
						FQDN:         fqdn,
					}
				}

				pc.serviceToMTLSMode[NewTargetServiceWithPortNumber(fqdn, portNumber)] = modeAndResource
			} else if port.GetNumber() != 0 {
				pc.serviceToMTLSMode[NewTargetServiceWithPortNumber(fqdn, port.GetNumber())] = modeAndResource
			} else {
				// Unhandled case!
				return fmt.Errorf("policy has a port with no name/number for target %s", target.Name)
			}
		}
	}

	return nil
}

// IsServiceMTLSEnforced returns true if a service requires incoming
// connections to use MTLS, or false if MTLS is not a hard-requirement (e.g.
// mode is permissive, peerIsOptional is true, etc). Only call this after adding
// all policy resources in effect via AddPolicy or AddMeshPolicy.
func (pc *PolicyChecker) IsServiceMTLSEnforced(w TargetService) (ModeAndResource, error) {
	// TODO support understanding port name -> port number mappings
	var modeAndResource ModeAndResource
	modeAndResource = pc.serviceToMTLSMode[w]
	if modeAndResource.Resource != nil {
		return modeAndResource, nil
	}

	// Try checking if its enforced on any ports
	serviceNoPort := NewTargetService(w.FQDN())
	modeAndResource = pc.serviceToMTLSMode[serviceNoPort]
	if modeAndResource.Resource != nil {
		return modeAndResource, nil
	}

	// Check if enforced on namespace
	namespace, _ := util.GetResourceNameFromHost("", w.FQDN()).InterpretAsNamespaceAndName()
	if namespace == "" {
		return ModeAndResource{}, fmt.Errorf("unable to extract namespace from fqdn: %s", w.FQDN())
	}
	modeAndResource = pc.namespaceToMTLSMode[namespace]
	if modeAndResource.Resource != nil {
		return modeAndResource, nil
	}

	// Finally, defer to mesh level policy. No need to check for a nil resource
	// since the default value is correct.
	return pc.meshMTLSModeAndResource, nil
}

// TargetServices returns the list of services known to the policy checker (in
// no particular order).
func (pc *PolicyChecker) TargetServices() []TargetService {
	tss := make([]TargetService, len(pc.serviceToMTLSMode))
	i := 0
	for ts := range pc.serviceToMTLSMode {
		tss[i] = ts
		i++
	}

	return tss
}

// parsePolicyMTLSMode is a helper function to determine what mtls mode a Policy
// implies.
func parsePolicyMTLSMode(p *v1alpha1.Policy) (Mode, error) {
	for _, peer := range p.Peers {
		mtlsParams, ok := peer.Params.(*v1alpha1.PeerAuthenticationMethod_Mtls)
		if !ok {
			// Only looking for mtls methods
			continue
		}

		var mode Mode
		if mtlsParams.Mtls == nil {
			mode = ModeStrict
		} else {
			switch mtlsParams.Mtls.GetMode() {
			case v1alpha1.MutualTls_PERMISSIVE:
				mode = ModePermissive
			case v1alpha1.MutualTls_STRICT:
				mode = ModeStrict
			default:
				// Shouldn't happen!
				return mode, fmt.Errorf("unknown MTLS mode when analyzing policy: %s", mtlsParams.Mtls.GetMode().String())
			}
		}
		// Now check for modifiers that might downgrade strict to permissive.
		if mode == ModeStrict && mtlsParams.Mtls != nil && mtlsParams.Mtls.AllowTls {
			mode = ModePermissive
		}
		if mode == ModeStrict && p.PeerIsOptional {
			mode = ModePermissive
		}
		return mode, nil
	}

	// No MTLS configuration found
	return ModePlaintext, nil
}
<|endoftext|>


<commit_before>// Package geoloader provides the interface between manager and dataset handling
// packages (geolite2 and legacy). manager only depends on geoloader and api.
// geoloader only depends on geolite2, legacy and api.
package geoloader

import (
	"context"
	"errors"
	"log"
	"regexp"
	"runtime"
	"sync"
	"time"

	"cloud.google.com/go/storage"
	"github.com/m-lab/annotation-service/api"
	"github.com/m-lab/annotation-service/geolite2"
	"github.com/m-lab/annotation-service/legacy"
	"github.com/m-lab/annotation-service/metrics"
	"google.golang.org/api/iterator"
)

var (
	// ErrNoLoader is returned if nil is passed for loader parameter.
	ErrNoLoader = errors.New("No loader provided")
)

// PopulateLatestData will search to the latest Geolite2 files
// available in GCS and will use them to create a new GeoDataset which
// it will place into the global scope as the latest version. It will
// do so safely with use of the currentDataMutex RW mutex. It it
// encounters an error, it will halt the program.
func GetLatestData() api.Annotator {
	data, err := geolite2.LoadLatestGeolite2File()
	if err != nil {
		log.Fatal(err)
	}
	return data
}

func ArchivedLoader(filename string) (api.Annotator, error) {
	if GeoLite2Regex.MatchString(filename) {
		return geolite2.LoadGeoLite2Dataset(filename, api.MaxmindBucketName)
	} else {
		return legacy.LoadLegacyDataset(filename, api.MaxmindBucketName)
	}
}

/*****************************************************************************
*                          LoadAll... functions                              *
*****************************************************************************/

// Returns the normal iterator for objects in the appropriate GCS bucket.
func bucketIterator() (*storage.ObjectIterator, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	prospectiveFiles := client.Bucket(api.MaxmindBucketName).Objects(ctx, &storage.Query{Prefix: api.MaxmindPrefix})
	return prospectiveFiles, nil
}

// LoadAll loads all datasets from the source that match the filter.
func LoadAll(
	source *storage.ObjectIterator,
	filter func(file *storage.ObjectAttrs) error,
	loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {

	result := make([]api.Annotator, 0, 100)
	resultLock := sync.Mutex{}
	wg := sync.WaitGroup{}

	for file, err := source.Next(); err != iterator.Done; file, err = source.Next() {
		// TODO - should we retry here?
		if err != nil {
			return nil, err
		}
		if file == nil {
			log.Println("file is nil", err)
			continue
		}
		if filter != nil && filter(file) != nil {
			continue
		}
		_, _, callerLine, _ := runtime.Caller(1)
		log.Println("Loading", file.Name, "from line", callerLine)
		wg.Add(1)
		go func(file *storage.ObjectAttrs) {
			defer wg.Done()
			ann, err := loader(file)
			if err != nil {
				log.Println("Retrying", file.Name, "after", err)
				ann, err = loader(file)
				if err != nil {
					log.Println("Failed trying to load", file.Name, "with", err)
					return
				}
			}
			resultLock.Lock()
			result = append(result, ann)
			resultLock.Unlock()
			metrics.DatasetCount.Inc()
			log.Println("Loaded", file.Name)
		}(file)
	}
	wg.Wait()
	return result, nil
}

func filter(file *storage.ObjectAttrs, r *regexp.Regexp, before time.Time) error {
	if !before.Equal(time.Time{}) {
		fileDate, err := api.ExtractDateFromFilename(file.Name)
		if err != nil {
			return err
		}
		if !fileDate.Before(GeoLite2StartDate) {
			return errors.New("After cutoff date") // TODO
		}
	}

	if !r.MatchString(file.Name) {
		return errors.New("Doesn't match") // TODO
	}

	return nil
}

// LoadAllLegacyV4 loads all v4 legacy datasets from the appropriate GCS bucket.
// The loader is injected, to allow for efficient unit testing.
func LoadAllLegacyV4(loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {
	if loader == nil {
		return nil, ErrNoLoader
	}
	it, err := bucketIterator()
	if err != nil {
		return nil, err
	}
	return LoadAll(it,
		func(file *storage.ObjectAttrs) error {
			return filter(file, GeoLegacyRegex, GeoLite2StartDate)
		},
		loader)
}

// LoadAllLegacyV6 loads all v6 legacy datasets from the appropriate GCS bucket.
// The loader is injected, to allow for efficient unit testing.
func LoadAllLegacyV6(loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {
	if loader == nil {
		return nil, ErrNoLoader
	}
	it, err := bucketIterator()
	if err != nil {
		return nil, err
	}
	return LoadAll(it,
		func(file *storage.ObjectAttrs) error {
			return filter(file, GeoLegacyv6Regex, GeoLite2StartDate)
		},
		loader)
}

// LoadGeolite2 loads a dataset from a GCS object.
func LoadGeolite2(file *storage.ObjectAttrs) (api.Annotator, error) {
	return geolite2.LoadGeoLite2Dataset(file.Name, file.Bucket)
}

// LoadAllGeolite2 loads all geolite2 datasets from the appropriate GCS bucket.
// The loader is injected, to allow for efficient unit testing.
func LoadAllGeolite2(loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {
	if loader == nil {
		return nil, ErrNoLoader
	}
	it, err := bucketIterator()
	if err != nil {
		return nil, err
	}
	return LoadAll(it,
		func(file *storage.ObjectAttrs) error {
			return filter(file, GeoLite2Regex, time.Time{})
		},
		loader)
}
<commit_msg>minor fixes<commit_after>// Package geoloader provides the interface between manager and dataset handling
// packages (geolite2 and legacy). manager only depends on geoloader and api.
// geoloader only depends on geolite2, legacy and api.
package geoloader

import (
	"context"
	"errors"
	"log"
	"regexp"
	"runtime"
	"sync"
	"time"

	"cloud.google.com/go/storage"
	"github.com/m-lab/annotation-service/api"
	"github.com/m-lab/annotation-service/geolite2"
	"github.com/m-lab/annotation-service/legacy"
	"github.com/m-lab/annotation-service/metrics"
	"google.golang.org/api/iterator"
)

var (
	// ErrNoLoader is returned if nil is passed for loader parameter.
	ErrNoLoader = errors.New("No loader provided")

	// ErrAfterLegacyCutoff is returned for legacy files beyond the cutoff date.
	errAfterLegacyCutoff = errors.New("After cutoff date")
	// ErrNoMatch is returned (internally) when filename does not match regexp.
	errNoMatch = errors.New("Doesn't match") // TODO
)

// PopulateLatestData will search to the latest Geolite2 files
// available in GCS and will use them to create a new GeoDataset which
// it will place into the global scope as the latest version. It will
// do so safely with use of the currentDataMutex RW mutex. It it
// encounters an error, it will halt the program.
func GetLatestData() api.Annotator {
	data, err := geolite2.LoadLatestGeolite2File()
	if err != nil {
		log.Fatal(err)
	}
	return data
}

func ArchivedLoader(filename string) (api.Annotator, error) {
	if GeoLite2Regex.MatchString(filename) {
		return geolite2.LoadGeoLite2Dataset(filename, api.MaxmindBucketName)
	} else {
		return legacy.LoadLegacyDataset(filename, api.MaxmindBucketName)
	}
}

/*****************************************************************************
*                          LoadAll... functions                              *
*****************************************************************************/

// Returns the normal iterator for objects in the appropriate GCS bucket.
func bucketIterator() (*storage.ObjectIterator, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	prospectiveFiles := client.Bucket(api.MaxmindBucketName).Objects(ctx, &storage.Query{Prefix: api.MaxmindPrefix})
	return prospectiveFiles, nil
}

// LoadAll loads all datasets from the source that match the filter.
func LoadAll(
	source *storage.ObjectIterator,
	filter func(file *storage.ObjectAttrs) error,
	loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {

	result := make([]api.Annotator, 0, 100)
	resultLock := sync.Mutex{}
	wg := sync.WaitGroup{}

	for file, err := source.Next(); err != iterator.Done; file, err = source.Next() {
		// TODO - should we retry here?
		if err != nil {
			return nil, err
		}
		if file == nil {
			log.Println("file is nil", err)
			continue
		}
		if filter != nil && filter(file) != nil {
			continue
		}
		_, _, callerLine, _ := runtime.Caller(1)
		log.Println("Loading", file.Name, "from line", callerLine)
		wg.Add(1)
		go func(file *storage.ObjectAttrs) {
			defer wg.Done()
			ann, err := loader(file)
			if err != nil {
				log.Println("Retrying", file.Name, "after", err)
				ann, err = loader(file)
				if err != nil {
					log.Println("Failed trying to load", file.Name, "with", err)
					return
				}
			}
			resultLock.Lock()
			result = append(result, ann)
			resultLock.Unlock()
			metrics.DatasetCount.Inc()
			log.Println("Loaded", file.Name)
		}(file)
	}
	wg.Wait()
	return result, nil
}

// filter is used to create filter functions for the loaders.
func filter(file *storage.ObjectAttrs, r *regexp.Regexp, before time.Time) error {
	if !before.Equal(time.Time{}) {
		fileDate, err := api.ExtractDateFromFilename(file.Name)
		if err != nil {
			return err
		}
		if !fileDate.Before(GeoLite2StartDate) {
			return errAfterLegacyCutoff
		}
	}

	if !r.MatchString(file.Name) {
		return errNoMatch
	}

	return nil
}

// LoadAllLegacyV4 loads all v4 legacy datasets from the appropriate GCS bucket.
// The loader is injected, to allow for efficient unit testing.
func LoadAllLegacyV4(loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {
	if loader == nil {
		return nil, ErrNoLoader
	}
	it, err := bucketIterator()
	if err != nil {
		return nil, err
	}
	return LoadAll(it,
		func(file *storage.ObjectAttrs) error {
			return filter(file, GeoLegacyRegex, GeoLite2StartDate)
		},
		loader)
}

// LoadAllLegacyV6 loads all v6 legacy datasets from the appropriate GCS bucket.
// The loader is injected, to allow for efficient unit testing.
func LoadAllLegacyV6(loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {
	if loader == nil {
		return nil, ErrNoLoader
	}
	it, err := bucketIterator()
	if err != nil {
		return nil, err
	}
	return LoadAll(it,
		func(file *storage.ObjectAttrs) error {
			return filter(file, GeoLegacyv6Regex, GeoLite2StartDate)
		},
		loader)
}

// LoadGeolite2 loads a dataset from a GCS object.
func LoadGeolite2(file *storage.ObjectAttrs) (api.Annotator, error) {
	return geolite2.LoadGeoLite2Dataset(file.Name, file.Bucket)
}

// LoadAllGeolite2 loads all geolite2 datasets from the appropriate GCS bucket.
// The loader is injected, to allow for efficient unit testing.
func LoadAllGeolite2(loader func(*storage.ObjectAttrs) (api.Annotator, error)) ([]api.Annotator, error) {
	if loader == nil {
		return nil, ErrNoLoader
	}
	it, err := bucketIterator()
	if err != nil {
		return nil, err
	}
	return LoadAll(it,
		func(file *storage.ObjectAttrs) error {
			return filter(file, GeoLite2Regex, time.Time{})
		},
		loader)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/keyserver/hashchain"
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/msg/padding"
	"github.com/mutecomm/mute/uid"
	"github.com/mutecomm/mute/uid/identity"
	"github.com/mutecomm/mute/uid/length"
	"golang.org/x/crypto/nacl/box"
)

// lengthEncryptedHeader defines the length of an encrypted header.
// This must always be the same in all messages! 5795 + 1373 = 7168.
const lengthEncryptedHeader = 7168

// Some wiggle room which can be taken out of padding if the need arises.
const wiggleRoom = 1373

type header struct {
	Ciphersuite                 string
	RecipientPubHash            string
	RecipientTempHash           string
	SenderSessionPub            uid.KeyEntry
	SenderIdentity              string // identity of sender
	SenderIdentityPubHash       string
	SenderIdentityPub           uid.KeyEntry
	NextSenderSessionPub        *uid.KeyEntry // optional
	NextRecipientSessionPubSeen *uid.KeyEntry // only if seen
	NymAddress                  string
	MaxDelay                    uint64
	SenderSessionCount          uint64
	SenderMessageCount          uint64
	SenderUID                   string // complete UID message in JSON
	SenderLastKeychainHash      string
	Status                      uint8
	Padding                     string
}

type headerPacket struct {
	Nonce                 [24]byte // for NaCL
	LengthEncryptedHeader uint16
	EncryptedHeader       []byte
}

func newHeader(
	sender, recipient *uid.Message,
	recipientTemp, senderSession, nextSenderSessionPub,
	nextRecipientSessionPubSeen *uid.KeyEntry,
	senderLastKeychainHash string,
	rand io.Reader,
) (*header, error) {
	if len(senderLastKeychainHash) != hashchain.EntryBase64Len {
		return nil, log.Errorf("msg: last hashchain entry '%s' does not have base64 length %d (but %d)",
			senderLastKeychainHash, hashchain.EntryBase64Len, len(senderLastKeychainHash))
	}
	h := &header{
		Ciphersuite:                 uid.DefaultCiphersuite, // at the moment we only support one ciphersuite
		RecipientPubHash:            recipient.PubHash(),
		RecipientTempHash:           recipientTemp.HASH,
		SenderIdentity:              sender.Identity(),
		SenderSessionPub:            *senderSession,
		SenderIdentityPubHash:       sender.PubHash(),
		SenderIdentityPub:           *sender.PubKey(),
		NextSenderSessionPub:        nextSenderSessionPub,
		NextRecipientSessionPubSeen: nextRecipientSessionPubSeen,
		NymAddress:                  sender.UIDContent.NYMADDRESS, // TODO: set the correct nymaddress!
		MaxDelay:                    0,                            // TODO
		SenderSessionCount:          0,                            // TODO
		SenderMessageCount:          0,                            // TODO
		SenderUID:                   string(sender.JSON()),
		SenderLastKeychainHash:      senderLastKeychainHash,
		Status:                      0,  // TODO
		Padding:                     "", // is set below
	}

	// calculate padding length
	padLen := wiggleRoom
	// pad sender identity
	if len(h.SenderIdentity) > identity.MaxLen {
		return nil, log.Error("msg: sender identity is too long")
	}
	padLen += identity.MaxLen - len(h.SenderIdentity)
	// pad nextSenderSessionPub
	if nextSenderSessionPub == nil {
		padLen += length.KeyEntryECDHE25519 - length.Nil
	}
	// pad nextRecipientSessionPubSeen
	if nextRecipientSessionPubSeen == nil {
		padLen += length.KeyEntryECDHE25519 - length.Nil
	}
	// pad nym address
	if len(h.NymAddress) > length.MaxNymAddress {
		return nil, log.Error("msg: nym address is too long")
	}
	padLen += length.MaxNymAddress - len(h.NymAddress)
	// pad sender UIDMessage
	if len(h.SenderUID) > length.MaxUIDMessage {
		return nil, log.Error("msg: sender UIDMesssage is too long")
	}
	padLen += length.MaxUIDMessage - len(h.SenderUID)
	// generate padding
	randLen := padLen/2 + padLen%2
	pad, err := padding.Generate(randLen, cipher.RandReader)
	if err != nil {
		return nil, err
	}
	// set padding
	p := hex.EncodeToString(pad)
	if padLen%2 == 1 {
		p = p[:len(p)-1]
	}
	h.Padding = p
	return h, nil
}

func newHeaderPacket(h *header, recipientIdentityPub, senderHeaderPriv *[32]byte, rand io.Reader) (*headerPacket, error) {
	var hp headerPacket
	jsn, err := json.Marshal(h)
	if err != nil {
		return nil, err
	}

	// TODO: show encrypted header contents
	// jsn2, _ := json.MarshalIndent(h, "", "  ")
	// fmt.Fprintln(os.Stderr, string(jsn2))

	if _, err := io.ReadFull(rand, hp.Nonce[:]); err != nil {
		return nil, log.Error(err)
	}
	hp.EncryptedHeader = box.Seal(hp.EncryptedHeader, jsn, &hp.Nonce, recipientIdentityPub, senderHeaderPriv)
	hp.LengthEncryptedHeader = uint16(len(hp.EncryptedHeader))
	if hp.LengthEncryptedHeader != lengthEncryptedHeader {
		return nil, log.Errorf("msg: encrypted header has wrong length (%d != %d)",
			hp.LengthEncryptedHeader, lengthEncryptedHeader)
	}
	return &hp, nil
}

func (hp *headerPacket) write(w io.Writer) error {
	log.Debugf("hp.Nonce: %s", base64.Encode(hp.Nonce[:]))
	if _, err := w.Write(hp.Nonce[:]); err != nil {
		return log.Error(err)
	}
	log.Debugf("hp.LengthEncryptedHeader: %d", hp.LengthEncryptedHeader)
	if err := binary.Write(w, binary.BigEndian, hp.LengthEncryptedHeader); err != nil {
		return log.Error(err)
	}
	if _, err := w.Write(hp.EncryptedHeader); err != nil {
		return log.Error(err)
	}
	return nil
}

func readHeader(
	senderHeaderPub *[32]byte,
	recipientIdentities []*uid.KeyEntry,
	r io.Reader,
) (int, *uid.KeyEntry, *header, error) {
	var hp headerPacket
	// read nonce
	if _, err := io.ReadFull(r, hp.Nonce[:]); err != nil {
		return 0, nil, nil, log.Error(err)
	}
	log.Debugf("hp.Nonce: %s", base64.Encode(hp.Nonce[:]))
	// read length of encrypted header
	if err := binary.Read(r, binary.BigEndian, &hp.LengthEncryptedHeader); err != nil {
		return 0, nil, nil, log.Error(err)
	}
	log.Debugf("hp.LengthEncryptedHeader: %d", hp.LengthEncryptedHeader)
	// read encrypted header
	hp.EncryptedHeader = make([]byte, hp.LengthEncryptedHeader)
	if _, err := io.ReadFull(r, hp.EncryptedHeader); err != nil {
		return 0, nil, nil, log.Error(err)
	}
	// try to decrypt header
	var jsn []byte
	var suc bool
	var recipientID *uid.KeyEntry
	var i int
	for idx, ke := range recipientIdentities {
		jsn, suc = box.Open(jsn, hp.EncryptedHeader, &hp.Nonce, senderHeaderPub, ke.PrivateKey32())
		if suc {
			i = idx
			recipientID = ke
			break
		}
	}
	if !suc {
		return 0, nil, nil, log.Error("msg: could not find key to decrypt header")
	}
	var h header
	if err := json.Unmarshal(jsn, &h); err != nil {
		return 0, nil, nil, err
	}

	// TODO: show encrypted header contents
	// jsn2, _ := json.MarshalIndent(h, "", "  ")
	// fmt.Fprintln(os.Stderr, string(jsn2))

	return i, recipientID, &h, nil
}
<commit_msg>msg: pad integers in header<commit_after>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/keyserver/hashchain"
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/msg/padding"
	"github.com/mutecomm/mute/uid"
	"github.com/mutecomm/mute/uid/identity"
	"github.com/mutecomm/mute/uid/length"
	"github.com/mutecomm/mute/util/digits"
	"golang.org/x/crypto/nacl/box"
)

// lengthEncryptedHeader defines the length of an encrypted header.
// This must always be the same in all messages! 5852 + 1316 = 7168.
const lengthEncryptedHeader = 7168

// Some wiggle room which can be taken out of padding if the need arises.
const wiggleRoom = 1316

type header struct {
	Ciphersuite                 string
	RecipientPubHash            string
	RecipientTempHash           string
	SenderSessionPub            uid.KeyEntry
	SenderIdentity              string // identity of sender
	SenderIdentityPubHash       string
	SenderIdentityPub           uid.KeyEntry
	NextSenderSessionPub        *uid.KeyEntry // optional
	NextRecipientSessionPubSeen *uid.KeyEntry // only if seen
	NymAddress                  string
	MaxDelay                    uint64
	SenderSessionCount          uint64
	SenderMessageCount          uint64
	SenderUID                   string // complete UID message in JSON
	SenderLastKeychainHash      string
	Status                      uint8 // always a single digit in JSON!
	Padding                     string
}

type headerPacket struct {
	Nonce                 [24]byte // for NaCL
	LengthEncryptedHeader uint16
	EncryptedHeader       []byte
}

func newHeader(
	sender, recipient *uid.Message,
	recipientTemp, senderSession, nextSenderSessionPub,
	nextRecipientSessionPubSeen *uid.KeyEntry,
	senderLastKeychainHash string,
	rand io.Reader,
) (*header, error) {
	if len(senderLastKeychainHash) != hashchain.EntryBase64Len {
		return nil, log.Errorf("msg: last hashchain entry '%s' does not have base64 length %d (but %d)",
			senderLastKeychainHash, hashchain.EntryBase64Len, len(senderLastKeychainHash))
	}
	h := &header{
		Ciphersuite:                 uid.DefaultCiphersuite, // at the moment we only support one ciphersuite
		RecipientPubHash:            recipient.PubHash(),
		RecipientTempHash:           recipientTemp.HASH,
		SenderIdentity:              sender.Identity(),
		SenderSessionPub:            *senderSession,
		SenderIdentityPubHash:       sender.PubHash(),
		SenderIdentityPub:           *sender.PubKey(),
		NextSenderSessionPub:        nextSenderSessionPub,
		NextRecipientSessionPubSeen: nextRecipientSessionPubSeen,
		NymAddress:                  sender.UIDContent.NYMADDRESS, // TODO: set the correct nymaddress!
		MaxDelay:                    0,                            // TODO
		SenderSessionCount:          0,                            // TODO
		SenderMessageCount:          0,                            // TODO
		SenderUID:                   string(sender.JSON()),
		SenderLastKeychainHash:      senderLastKeychainHash,
		Status:                      0,  // TODO
		Padding:                     "", // is set below
	}

	// calculate padding length
	padLen := wiggleRoom
	// pad sender identity
	if len(h.SenderIdentity) > identity.MaxLen {
		return nil, log.Error("msg: sender identity is too long")
	}
	padLen += identity.MaxLen - len(h.SenderIdentity)
	// pad nextSenderSessionPub
	if nextSenderSessionPub == nil {
		padLen += length.KeyEntryECDHE25519 - length.Nil
	}
	// pad nextRecipientSessionPubSeen
	if nextRecipientSessionPubSeen == nil {
		padLen += length.KeyEntryECDHE25519 - length.Nil
	}
	// pad nym address
	if len(h.NymAddress) > length.MaxNymAddress {
		return nil, log.Error("msg: nym address is too long")
	}
	padLen += length.MaxNymAddress - len(h.NymAddress)
	// pad integers
	padLen += 20 - digits.Count(h.MaxDelay)
	padLen += 20 - digits.Count(h.SenderSessionCount)
	padLen += 20 - digits.Count(h.SenderMessageCount)
	// pad sender UIDMessage
	if len(h.SenderUID) > length.MaxUIDMessage {
		return nil, log.Error("msg: sender UIDMesssage is too long")
	}
	padLen += length.MaxUIDMessage - len(h.SenderUID)
	// generate padding
	randLen := padLen/2 + padLen%2
	pad, err := padding.Generate(randLen, cipher.RandReader)
	if err != nil {
		return nil, err
	}
	// set padding
	p := hex.EncodeToString(pad)
	if padLen%2 == 1 {
		p = p[:len(p)-1]
	}
	h.Padding = p
	return h, nil
}

func newHeaderPacket(h *header, recipientIdentityPub, senderHeaderPriv *[32]byte, rand io.Reader) (*headerPacket, error) {
	var hp headerPacket
	jsn, err := json.Marshal(h)
	if err != nil {
		return nil, err
	}

	// TODO: show encrypted header contents
	// jsn2, _ := json.MarshalIndent(h, "", "  ")
	// fmt.Fprintln(os.Stderr, string(jsn2))

	if _, err := io.ReadFull(rand, hp.Nonce[:]); err != nil {
		return nil, log.Error(err)
	}
	hp.EncryptedHeader = box.Seal(hp.EncryptedHeader, jsn, &hp.Nonce, recipientIdentityPub, senderHeaderPriv)
	hp.LengthEncryptedHeader = uint16(len(hp.EncryptedHeader))
	if hp.LengthEncryptedHeader != lengthEncryptedHeader {
		return nil, log.Errorf("msg: encrypted header has wrong length (%d != %d)",
			hp.LengthEncryptedHeader, lengthEncryptedHeader)
	}
	return &hp, nil
}

func (hp *headerPacket) write(w io.Writer) error {
	log.Debugf("hp.Nonce: %s", base64.Encode(hp.Nonce[:]))
	if _, err := w.Write(hp.Nonce[:]); err != nil {
		return log.Error(err)
	}
	log.Debugf("hp.LengthEncryptedHeader: %d", hp.LengthEncryptedHeader)
	if err := binary.Write(w, binary.BigEndian, hp.LengthEncryptedHeader); err != nil {
		return log.Error(err)
	}
	if _, err := w.Write(hp.EncryptedHeader); err != nil {
		return log.Error(err)
	}
	return nil
}

func readHeader(
	senderHeaderPub *[32]byte,
	recipientIdentities []*uid.KeyEntry,
	r io.Reader,
) (int, *uid.KeyEntry, *header, error) {
	var hp headerPacket
	// read nonce
	if _, err := io.ReadFull(r, hp.Nonce[:]); err != nil {
		return 0, nil, nil, log.Error(err)
	}
	log.Debugf("hp.Nonce: %s", base64.Encode(hp.Nonce[:]))
	// read length of encrypted header
	if err := binary.Read(r, binary.BigEndian, &hp.LengthEncryptedHeader); err != nil {
		return 0, nil, nil, log.Error(err)
	}
	log.Debugf("hp.LengthEncryptedHeader: %d", hp.LengthEncryptedHeader)
	// read encrypted header
	hp.EncryptedHeader = make([]byte, hp.LengthEncryptedHeader)
	if _, err := io.ReadFull(r, hp.EncryptedHeader); err != nil {
		return 0, nil, nil, log.Error(err)
	}
	// try to decrypt header
	var jsn []byte
	var suc bool
	var recipientID *uid.KeyEntry
	var i int
	for idx, ke := range recipientIdentities {
		jsn, suc = box.Open(jsn, hp.EncryptedHeader, &hp.Nonce, senderHeaderPub, ke.PrivateKey32())
		if suc {
			i = idx
			recipientID = ke
			break
		}
	}
	if !suc {
		return 0, nil, nil, log.Error("msg: could not find key to decrypt header")
	}
	var h header
	if err := json.Unmarshal(jsn, &h); err != nil {
		return 0, nil, nil, err
	}

	// TODO: show encrypted header contents
	// jsn2, _ := json.MarshalIndent(h, "", "  ")
	// fmt.Fprintln(os.Stderr, string(jsn2))

	return i, recipientID, &h, nil
}
<|endoftext|>


<commit_before>package mustache

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"testing"
)

var enabledTests = map[string]map[string]bool{
	"comments.json": map[string]bool{
		"Inline":                           true,
		"Multiline":                        true,
		"Standalone":                       false,
		"Indented Standalone":              false,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
		"Multiline Standalone":             false,
		"Indented Multiline Standalone":    false,
		"Indented Inline":                  true,
		"Surrounding Whitespace":           true,
	},
	"delimiters.json": map[string]bool{
		"Pair Behavior":                    true,
		"Special Characters":               true,
		"Sections":                         false,
		"Inverted Sections":                false,
		"Partial Inheritence":              false,
		"Post-Partial Behavior":            false,
		"Outlying Whitespace (Inline)":     true,
		"Standalone Tag":                   false,
		"Indented Standalone Tag":          false,
		"Pair with Padding":                true,
		"Surrounding Whitespace":           true,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
	},
	"interpolation.json": map[string]bool{
		"No Interpolation":                             true,
		"Basic Interpolation":                          true,
		"HTML Escaping":                                false,
		"Triple Mustache":                              true,
		"Ampersand":                                    false,
		"Basic Integer Interpolation":                  true,
		"Triple Mustache Integer Interpolation":        true,
		"Ampersand Integer Interpolation":              false,
		"Basic Decimal Interpolation":                  true,
		"Triple Mustache Decimal Interpolation":        true,
		"Ampersand Decimal Interpolation":              false,
		"Basic Context Miss Interpolation":             true,
		"Triple Mustache Context Miss Interpolation":   true,
		"Ampersand Context Miss Interpolation":         true,
		"Dotted Names - Basic Interpolation":           true,
		"Dotted Names - Triple Mustache Interpolation": true,
		"Dotted Names - Ampersand Interpolation":       false,
		"Dotted Names - Arbitrary Depth":               true,
		"Dotted Names - Broken Chains":                 true,
		"Dotted Names - Broken Chain Resolution":       true,
		"Dotted Names - Initial Resolution":            true,
		"Interpolation - Surrounding Whitespace":       true,
		"Triple Mustache - Surrounding Whitespace":     true,
		"Ampersand - Surrounding Whitespace":           false,
		"Interpolation - Standalone":                   true,
		"Triple Mustache - Standalone":                 true,
		"Ampersand - Standalone":                       false,
		"Interpolation With Padding":                   true,
		"Triple Mustache With Padding":                 false,
		"Ampersand With Padding":                       false,
	},
	"inverted.json": map[string]bool{
		"Falsey":                           true,
		"Truthy":                           true,
		"Context":                          true,
		"List":                             true,
		"Empty List":                       true,
		"Doubled":                          false,
		"Nested (Falsey)":                  true,
		"Nested (Truthy)":                  true,
		"Context Misses":                   true,
		"Dotted Names - Truthy":            true,
		"Dotted Names - Falsey":            true,
		"Internal Whitespace":              true,
		"Indented Inline Sections":         true,
		"Standalone Lines":                 false,
		"Standalone Indented Lines":        false,
		"Padding":                          true,
		"Dotted Names - Broken Chains":     true,
		"Surrounding Whitespace":           true,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
	},
	"partials.json": map[string]bool{
		"Basic Behavior":                   true,
		"Failed Lookup":                    true,
		"Context":                          true,
		"Recursion":                        true,
		"Surrounding Whitespace":           true,
		"Inline Indentation":               true,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
		"Standalone Indentation":           false,
		"Padding Whitespace":               true,
	},
	"sections.json": map[string]bool{
		"Truthy":                 true,
		"Falsey":                 true,
		"Context":                true,
		"Deeply Nested Contexts": false,
		"List":                             true,
		"Empty List":                       true,
		"Doubled":                          false,
		"Nested (Truthy)":                  true,
		"Nested (Falsey)":                  true,
		"Context Misses":                   true,
		"Implicit Iterator - String":       true,
		"Implicit Iterator - Integer":      true,
		"Implicit Iterator - Decimal":      true,
		"Implicit Iterator - Array":        true,
		"Dotted Names - Truthy":            true,
		"Dotted Names - Falsey":            true,
		"Dotted Names - Broken Chains":     true,
		"Surrounding Whitespace":           true,
		"Internal Whitespace":              true,
		"Indented Inline Sections":         true,
		"Standalone Lines":                 false,
		"Indented Standalone Lines":        false,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
		"Padding":                          true,
	},
	"~lambdas.json": nil, // not implemented
}

type specTest struct {
	Name        string            `json:"name"`
	Data        interface{}       `json:"data"`
	Expected    string            `json:"expected"`
	Template    string            `json:"template"`
	Description string            `json:"desc"`
	Partials    map[string]string `json:"partials"`
}

type specTestSuite struct {
	Tests []specTest `json:"tests"`
}

func TestSpec(t *testing.T) {
	root := filepath.Join(os.Getenv("PWD"), "spec", "specs")
	if _, err := os.Stat(root); err != nil {
		if os.IsNotExist(err) {
			t.Fatalf("Could not find the specs folder at %s, ensure the submodule exists by running 'git submodule update --init'", root)
		}
		t.Fatal(err)
	}

	paths, err := filepath.Glob(root + "/*.json")
	if err != nil {
		t.Fatal(err)
	}
	sort.Strings(paths)

	for _, path := range paths {
		_, file := filepath.Split(path)
		enabled, ok := enabledTests[file]
		if !ok {
			t.Errorf("Unexpected file %s, consider adding to enabledFiles", file)
			continue
		}
		if enabled == nil {
			continue
		}
		b, err := ioutil.ReadFile(path)
		if err != nil {
			t.Fatal(err)
		}
		var suite specTestSuite
		err = json.Unmarshal(b, &suite)
		if err != nil {
			t.Fatal(err)
		}
		for _, test := range suite.Tests {
			runTest(t, file, &test)
		}
	}
}

func runTest(t *testing.T, file string, test *specTest) {
	enabled, ok := enabledTests[file][test.Name]
	if !ok {
		t.Errorf("[%s %s]: Unexpected test, add to enabledTests", file, test.Name)
	}
	if !enabled {
		t.Logf("[%s %s]: Skipped", file, test.Name)
		return
	}

	var out string
	var err error
	if len(test.Partials) > 0 {
		out, err = RenderPartials(test.Template, StaticProvider(test.Partials), test.Data)
	} else {
		out, err = Render(test.Template, test.Data)
	}
	if err != nil {
		t.Errorf("[%s %s]: %s", file, test.Name, err.Error())
		return
	}
	if out != test.Expected {
		t.Errorf("[%s %s]: Expected %q, got %q", file, test.Name, test.Expected, out)
		return
	}

	t.Logf("[%s %s]: Passed", file, test.Name)
}
<commit_msg>Enable one more test<commit_after>package mustache

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"testing"
)

var enabledTests = map[string]map[string]bool{
	"comments.json": map[string]bool{
		"Inline":                           true,
		"Multiline":                        true,
		"Standalone":                       false,
		"Indented Standalone":              false,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
		"Multiline Standalone":             false,
		"Indented Multiline Standalone":    false,
		"Indented Inline":                  true,
		"Surrounding Whitespace":           true,
	},
	"delimiters.json": map[string]bool{
		"Pair Behavior":                    true,
		"Special Characters":               true,
		"Sections":                         false,
		"Inverted Sections":                false,
		"Partial Inheritence":              false,
		"Post-Partial Behavior":            true,
		"Outlying Whitespace (Inline)":     true,
		"Standalone Tag":                   false,
		"Indented Standalone Tag":          false,
		"Pair with Padding":                true,
		"Surrounding Whitespace":           true,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
	},
	"interpolation.json": map[string]bool{
		"No Interpolation":                             true,
		"Basic Interpolation":                          true,
		"HTML Escaping":                                false,
		"Triple Mustache":                              true,
		"Ampersand":                                    false,
		"Basic Integer Interpolation":                  true,
		"Triple Mustache Integer Interpolation":        true,
		"Ampersand Integer Interpolation":              false,
		"Basic Decimal Interpolation":                  true,
		"Triple Mustache Decimal Interpolation":        true,
		"Ampersand Decimal Interpolation":              false,
		"Basic Context Miss Interpolation":             true,
		"Triple Mustache Context Miss Interpolation":   true,
		"Ampersand Context Miss Interpolation":         true,
		"Dotted Names - Basic Interpolation":           true,
		"Dotted Names - Triple Mustache Interpolation": true,
		"Dotted Names - Ampersand Interpolation":       false,
		"Dotted Names - Arbitrary Depth":               true,
		"Dotted Names - Broken Chains":                 true,
		"Dotted Names - Broken Chain Resolution":       true,
		"Dotted Names - Initial Resolution":            true,
		"Interpolation - Surrounding Whitespace":       true,
		"Triple Mustache - Surrounding Whitespace":     true,
		"Ampersand - Surrounding Whitespace":           false,
		"Interpolation - Standalone":                   true,
		"Triple Mustache - Standalone":                 true,
		"Ampersand - Standalone":                       false,
		"Interpolation With Padding":                   true,
		"Triple Mustache With Padding":                 false,
		"Ampersand With Padding":                       false,
	},
	"inverted.json": map[string]bool{
		"Falsey":                           true,
		"Truthy":                           true,
		"Context":                          true,
		"List":                             true,
		"Empty List":                       true,
		"Doubled":                          false,
		"Nested (Falsey)":                  true,
		"Nested (Truthy)":                  true,
		"Context Misses":                   true,
		"Dotted Names - Truthy":            true,
		"Dotted Names - Falsey":            true,
		"Internal Whitespace":              true,
		"Indented Inline Sections":         true,
		"Standalone Lines":                 false,
		"Standalone Indented Lines":        false,
		"Padding":                          true,
		"Dotted Names - Broken Chains":     true,
		"Surrounding Whitespace":           true,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
	},
	"partials.json": map[string]bool{
		"Basic Behavior":                   true,
		"Failed Lookup":                    true,
		"Context":                          true,
		"Recursion":                        true,
		"Surrounding Whitespace":           true,
		"Inline Indentation":               true,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
		"Standalone Indentation":           false,
		"Padding Whitespace":               true,
	},
	"sections.json": map[string]bool{
		"Truthy":                 true,
		"Falsey":                 true,
		"Context":                true,
		"Deeply Nested Contexts": false,
		"List":                             true,
		"Empty List":                       true,
		"Doubled":                          false,
		"Nested (Truthy)":                  true,
		"Nested (Falsey)":                  true,
		"Context Misses":                   true,
		"Implicit Iterator - String":       true,
		"Implicit Iterator - Integer":      true,
		"Implicit Iterator - Decimal":      true,
		"Implicit Iterator - Array":        true,
		"Dotted Names - Truthy":            true,
		"Dotted Names - Falsey":            true,
		"Dotted Names - Broken Chains":     true,
		"Surrounding Whitespace":           true,
		"Internal Whitespace":              true,
		"Indented Inline Sections":         true,
		"Standalone Lines":                 false,
		"Indented Standalone Lines":        false,
		"Standalone Line Endings":          false,
		"Standalone Without Previous Line": false,
		"Standalone Without Newline":       false,
		"Padding":                          true,
	},
	"~lambdas.json": nil, // not implemented
}

type specTest struct {
	Name        string            `json:"name"`
	Data        interface{}       `json:"data"`
	Expected    string            `json:"expected"`
	Template    string            `json:"template"`
	Description string            `json:"desc"`
	Partials    map[string]string `json:"partials"`
}

type specTestSuite struct {
	Tests []specTest `json:"tests"`
}

func TestSpec(t *testing.T) {
	root := filepath.Join(os.Getenv("PWD"), "spec", "specs")
	if _, err := os.Stat(root); err != nil {
		if os.IsNotExist(err) {
			t.Fatalf("Could not find the specs folder at %s, ensure the submodule exists by running 'git submodule update --init'", root)
		}
		t.Fatal(err)
	}

	paths, err := filepath.Glob(root + "/*.json")
	if err != nil {
		t.Fatal(err)
	}
	sort.Strings(paths)

	for _, path := range paths {
		_, file := filepath.Split(path)
		enabled, ok := enabledTests[file]
		if !ok {
			t.Errorf("Unexpected file %s, consider adding to enabledFiles", file)
			continue
		}
		if enabled == nil {
			continue
		}
		b, err := ioutil.ReadFile(path)
		if err != nil {
			t.Fatal(err)
		}
		var suite specTestSuite
		err = json.Unmarshal(b, &suite)
		if err != nil {
			t.Fatal(err)
		}
		for _, test := range suite.Tests {
			runTest(t, file, &test)
		}
	}
}

func runTest(t *testing.T, file string, test *specTest) {
	enabled, ok := enabledTests[file][test.Name]
	if !ok {
		t.Errorf("[%s %s]: Unexpected test, add to enabledTests", file, test.Name)
	}
	if !enabled {
		t.Logf("[%s %s]: Skipped", file, test.Name)
		return
	}

	var out string
	var err error
	if len(test.Partials) > 0 {
		out, err = RenderPartials(test.Template, StaticProvider(test.Partials), test.Data)
	} else {
		out, err = Render(test.Template, test.Data)
	}
	if err != nil {
		t.Errorf("[%s %s]: %s", file, test.Name, err.Error())
		return
	}
	if out != test.Expected {
		t.Errorf("[%s %s]: Expected %q, got %q", file, test.Name, test.Expected, out)
		return
	}

	t.Logf("[%s %s]: Passed", file, test.Name)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/crackcomm/crawl"
	"github.com/golang/glog"
	"github.com/nsqio/go-nsq"
)

func main() {
	defer glog.Flush()

	// CRAWL_DEBUG environment variable turns on debug mode
	// crawler then can spit out logs using glog.V(3)
	var verbosity string
	if yes, _ := strconv.ParseBool(os.Getenv("CRAWL_DEBUG")); yes {
		verbosity = "-v=3"
	}

	// We are setting glog to log to stderr
	flag.CommandLine.Parse([]string{"-logtostderr", verbosity})

	app := cli.NewApp()
	app.Name = "crawl-schedule"
	app.HelpName = app.Name
	app.Version = "0.0.1"
	app.ArgsUsage = "<url>"
	app.Usage = "schedules a crawl request in nsq"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "nsq-addr",
			EnvVar: "NSQ_ADDR",
			Usage:  "nsq address",
		},
		cli.StringFlag{
			Name:   "topic",
			EnvVar: "TOPIC",
			Usage:  "crawl requests nsq topic",
			Value:  "crawl_requests",
		},
		cli.StringSliceFlag{
			Name:  "form-value",
			Usage: "form value in format (key=value)",
		},
		cli.StringSliceFlag{
			Name:  "metadata",
			Usage: "metadata value in format (key=value)",
		},
		cli.StringSliceFlag{
			Name:  "callback",
			Usage: "crawl request callbacks (at least one is required)",
		},
		cli.StringFlag{
			Name:  "referer",
			Usage: "crawl request referer",
		},
		cli.StringFlag{
			Name:  "method",
			Usage: "crawl request referer",
			Value: "GET",
		},
	}
	app.Before = func(c *cli.Context) error {
		var errs []string
		if len(c.String("topic")) == 0 {
			errs = append(errs, "Topic cannot be empty")
		}
		if len(c.String("nsq-addr")) == 0 {
			errs = append(errs, "At least one --nsq-addr is required")
		}
		if len(c.StringSlice("callback")) == 0 {
			errs = append(errs, "At least one --callback is required")
		}
		if len(c.Args()) != 1 {
			errs = append(errs, "At least one url is required in arguments.")
		}
		if len(errs) != 0 {
			return errors.New(strings.Join(errs, "\n"))
		}
		return nil
	}
	app.Action = func(c *cli.Context) {
		form, err := listToMap(c.StringSlice("form-value"))
		if err != nil {
			glog.Fatalf("Form values error: %v", err)
		}

		metadata, err := listToMap(c.StringSlice("metadata"))
		if err != nil {
			glog.Fatalf("Metadata error: %v", err)
		}

		request := &crawl.Request{
			URL:       strings.Trim(c.Args().First(), `"'`),
			Method:    c.String("method"),
			Referer:   c.String("referer"),
			Callbacks: c.StringSlice("callback"),
			Form:      form,
			Metadata:  metadata,
		}

		if glog.V(3) {
			body, _ := json.MarshalIndent(request, "", "  ")
			glog.Infof("Scheduling request: %s", body)
		}

		body, err := json.Marshal(request)
		if err != nil {
			glog.Fatalf("Error marshaling request to json: %v", err)
		}

		producer, err := nsq.NewProducer(c.String("nsq-addr"), nsq.NewConfig())
		if err != nil {
			glog.Fatalf("Error connecting to nsq: %v", err)
		}

		if err := producer.Publish(c.String("topic"), body); err != nil {
			glog.Fatalf("Publish error: %v", err)
		}

		glog.Info("Request scheduled")
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func listToMap(list []string) (result map[string]string, err error) {
	result = make(map[string]string)
	for _, keyValue := range list {
		i := strings.Index(keyValue, "=")
		if i <= 0 {
			return nil, fmt.Errorf("%q is not valid", keyValue)
		}
		key := keyValue[:i]
		value := keyValue[i+1:]
		result[key] = value
	}
	return
}
<commit_msg>nsq crawl-schedule: set nsq output buffer timeout to 0 + logging<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/crackcomm/crawl"
	"github.com/golang/glog"
	"github.com/nsqio/go-nsq"
)

func main() {
	defer glog.Flush()

	// CRAWL_DEBUG environment variable turns on debug mode
	// crawler then can spit out logs using glog.V(3)
	var verbosity string
	if yes, _ := strconv.ParseBool(os.Getenv("CRAWL_DEBUG")); yes {
		verbosity = "-v=3"
	}

	// We are setting glog to log to stderr
	flag.CommandLine.Parse([]string{"-logtostderr", verbosity})

	app := cli.NewApp()
	app.Name = "crawl-schedule"
	app.HelpName = app.Name
	app.Version = "0.0.1"
	app.ArgsUsage = "<url>"
	app.Usage = "schedules a crawl request in nsq"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "nsq-addr",
			EnvVar: "NSQ_ADDR",
			Usage:  "nsq address",
		},
		cli.StringFlag{
			Name:   "topic",
			EnvVar: "TOPIC",
			Usage:  "crawl requests nsq topic",
			Value:  "crawl_requests",
		},
		cli.StringSliceFlag{
			Name:  "form-value",
			Usage: "form value in format (key=value)",
		},
		cli.StringSliceFlag{
			Name:  "metadata",
			Usage: "metadata value in format (key=value)",
		},
		cli.StringSliceFlag{
			Name:  "callback",
			Usage: "crawl request callbacks (at least one is required)",
		},
		cli.StringFlag{
			Name:  "referer",
			Usage: "crawl request referer",
		},
		cli.StringFlag{
			Name:  "method",
			Usage: "crawl request referer",
			Value: "GET",
		},
	}
	app.Before = func(c *cli.Context) error {
		var errs []string
		if len(c.String("topic")) == 0 {
			errs = append(errs, "Topic cannot be empty")
		}
		if len(c.String("nsq-addr")) == 0 {
			errs = append(errs, "At least one --nsq-addr is required")
		}
		if len(c.StringSlice("callback")) == 0 {
			errs = append(errs, "At least one --callback is required")
		}
		if len(c.Args()) != 1 {
			errs = append(errs, "At least one url is required in arguments.")
		}
		if len(errs) != 0 {
			return errors.New(strings.Join(errs, "\n"))
		}
		return nil
	}
	app.Action = func(c *cli.Context) {
		form, err := listToMap(c.StringSlice("form-value"))
		if err != nil {
			glog.Fatalf("Form values error: %v", err)
		}

		metadata, err := listToMap(c.StringSlice("metadata"))
		if err != nil {
			glog.Fatalf("Metadata error: %v", err)
		}

		request := &crawl.Request{
			URL:       strings.Trim(c.Args().First(), `"'`),
			Method:    c.String("method"),
			Referer:   c.String("referer"),
			Callbacks: c.StringSlice("callback"),
			Form:      form,
			Metadata:  metadata,
		}

		if glog.V(3) {
			body, _ := json.MarshalIndent(request, "", "  ")
			glog.Infof("Scheduling request: %s", body)
		}

		body, err := json.Marshal(request)
		if err != nil {
			glog.Fatalf("Error marshaling request to json: %v", err)
		}

		cfg := nsq.NewConfig()
		cfg.OutputBufferTimeout = 0

		producer, err := nsq.NewProducer(c.String("nsq-addr"), cfg)
		if err != nil {
			glog.Fatalf("Error connecting to nsq: %v", err)
		}

		producer.SetLogger(log.New(os.Stdout, "[nsq]", 0), nsq.LogLevelError)

		glog.Infof("Publishing request to %q", c.String("topic"))

		if err := producer.Publish(c.String("topic"), body); err != nil {
			glog.Fatalf("Publish error: %v", err)
		}

		glog.Info("Request scheduled")
		producer.Stop()
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func listToMap(list []string) (result map[string]string, err error) {
	result = make(map[string]string)
	for _, keyValue := range list {
		i := strings.Index(keyValue, "=")
		if i <= 0 {
			return nil, fmt.Errorf("%q is not valid", keyValue)
		}
		key := keyValue[:i]
		value := keyValue[i+1:]
		result[key] = value
	}
	return
}
<|endoftext|>


<commit_before>/*
Package github_flavored_markdown provides a GitHub Flavored Markdown renderer
with fenced code block highlighting, clickable header anchor links.

The functionality should be equivalent to the GitHub Markdown API endpoint specified at
https://developer.github.com/v3/markdown/#render-a-markdown-document-in-raw-mode, except
the rendering is performed locally.

See example below for how to generate a complete HTML page, including CSS styles.
*/
package github_flavored_markdown

import (
	"bytes"
	"fmt"
	"html"
	"io"
	"regexp"
	"sort"
	"strings"
	"text/template"

	"github.com/microcosm-cc/bluemonday"
	"github.com/russross/blackfriday"
	"github.com/shurcooL/go/github_flavored_markdown/sanitized_anchor_name"
	"github.com/shurcooL/go/highlight_go"
	"github.com/shurcooL/go/u/u7"
	"github.com/sourcegraph/annotate"
	"github.com/sourcegraph/syntaxhighlight"
)

// Markdown renders GitHub Flavored Markdown text.
func Markdown(text []byte) []byte {
	htmlFlags := 0
	renderer := &renderer{Html: blackfriday.HtmlRenderer(htmlFlags, "", "").(*blackfriday.Html)}

	// Parser extensions for GitHub Flavored Markdown.
	extensions := 0
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS
	//extensions |= blackfriday.EXTENSION_HARD_LINE_BREAK

	unsanitized := blackfriday.Markdown(text, renderer, extensions)

	// GitHub Flavored Markdown-like sanitization policy.
	p := bluemonday.UGCPolicy()
	p.AllowAttrs("class").Matching(bluemonday.SpaceSeparatedTokens).OnElements("div", "span")
	p.AllowAttrs("class", "name").Matching(bluemonday.SpaceSeparatedTokens).OnElements("a")
	p.AllowAttrs("rel").Matching(regexp.MustCompile(`^nofollow$`)).OnElements("a")
	p.AllowAttrs("aria-hidden").Matching(regexp.MustCompile(`^true$`)).OnElements("a")
	p.AllowDataURIImages()

	return p.SanitizeBytes(unsanitized)
}

type renderer struct {
	*blackfriday.Html
}

// GitHub Flavored Markdown header with clickable and hidden anchor.
func (_ *renderer) Header(out *bytes.Buffer, text func() bool, level int, _ string) {
	marker := out.Len()
	doubleSpace(out)

	if !text() {
		out.Truncate(marker)
		return
	}

	textString := out.String()[marker:]
	out.Truncate(marker)

	anchorName := sanitized_anchor_name.Create(html.UnescapeString(textString))

	out.WriteString(fmt.Sprintf(`<h%d><a name="%s" class="anchor" href="#%s" rel="nofollow" aria-hidden="true"><span class="octicon octicon-link"></span></a>`, level, anchorName, anchorName))
	out.WriteString(textString)
	out.WriteString(fmt.Sprintf("</h%d>\n", level))
}

// TODO: Clean up and improve this code.
// GitHub Flavored Markdown fenced code block with highlighting.
func (_ *renderer) BlockCode(out *bytes.Buffer, text []byte, lang string) {
	doubleSpace(out)

	// parse out the language name
	count := 0
	for _, elt := range strings.Fields(lang) {
		if elt[0] == '.' {
			elt = elt[1:]
		}
		if len(elt) == 0 {
			continue
		}
		out.WriteString(`<div class="highlight highlight-`)
		attrEscape(out, []byte(elt))
		lang = elt
		out.WriteString(`"><pre>`)
		count++
		break
	}

	if count == 0 {
		out.WriteString("
")
	}

	if highlightedCode, ok := highlightCode(text, lang); ok {
		out.Write(highlightedCode)
	} else {
		attrEscape(out, text)
	}

	if count == 0 {
		out.WriteString("
\n")
	} else {
		out.WriteString("

\n")
	}
}

var gfmHtmlConfig = syntaxhighlight.HTMLConfig{
	String:        "s",
	Keyword:       "k",
	Comment:       "c",
	Type:          "n",
	Literal:       "o",
	Punctuation:   "p",
	Plaintext:     "n",
	Tag:           "tag",
	HTMLTag:       "htm",
	HTMLAttrName:  "atn",
	HTMLAttrValue: "atv",
	Decimal:       "m",
}

// TODO: Support highlighting for more languages.
func highlightCode(src []byte, lang string) (highlightedCode []byte, ok bool) {
	switch lang {
	case "Go":
		var buf bytes.Buffer
		err := highlight_go.Print(src, &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "Go-old":
		var buf bytes.Buffer
		err := syntaxhighlight.Print(syntaxhighlight.NewScanner(src), &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "diff":
		switch 2 {
		default:
			var buf bytes.Buffer
			err := u7.Print(u7.NewScanner(src), &buf)
			if err != nil {
				return nil, false
			}
			return buf.Bytes(), true
		case 1:
			lines := bytes.Split(src, []byte("\n"))
			return bytes.Join(lines, []byte("\n")), true
		case 2:
			anns, err := u7.Annotate(src)
			if err != nil {
				return nil, false
			}

			lines := bytes.Split(src, []byte("\n"))
			lineStarts := make([]int, len(lines))
			var offset int
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				lineStarts[lineIndex] = offset
				offset += len(lines[lineIndex]) + 1
			}

			lastDel, lastIns := -1, -1
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				var lineFirstChar byte
				if len(lines[lineIndex]) > 0 {
					lineFirstChar = lines[lineIndex][0]
				}
				switch lineFirstChar {
				case '+':
					if lastIns == -1 {
						lastIns = lineIndex
					}
				case '-':
					if lastDel == -1 {
						lastDel = lineIndex
					}
				default:
					if lastDel != -1 || lastIns != -1 {
						if lastDel == -1 {
							lastDel = lastIns
						} else if lastIns == -1 {
							lastIns = lineIndex
						}

						beginOffsetLeft := lineStarts[lastDel]
						endOffsetLeft := lineStarts[lastIns]
						beginOffsetRight := lineStarts[lastIns]
						endOffsetRight := lineStarts[lineIndex]

						anns = append(anns, &annotate.Annotation{Start: beginOffsetLeft, End: endOffsetLeft, Left: []byte(`<span class="gd input-block">`), Right: []byte(`</span>`), WantInner: 0})
						anns = append(anns, &annotate.Annotation{Start: beginOffsetRight, End: endOffsetRight, Left: []byte(`<span class="gi input-block">`), Right: []byte(`</span>`), WantInner: 0})

						if '@' != lineFirstChar {
							//leftContent := string(src[beginOffsetLeft:endOffsetLeft])
							//rightContent := string(src[beginOffsetRight:endOffsetRight])
							// This is needed to filter out the "-" and "+" at the beginning of each line from being highlighted.
							// TODO: Still not completely filtered out.
							leftContent := ""
							for line := lastDel; line < lastIns; line++ {
								leftContent += "\x00" + string(lines[line][1:]) + "\n"
							}
							rightContent := ""
							for line := lastIns; line < lineIndex; line++ {
								rightContent += "\x00" + string(lines[line][1:]) + "\n"
							}

							var sectionSegments [2][]*annotate.Annotation
							u7.HighlightedDiffFunc(leftContent, rightContent, &sectionSegments, [2]int{beginOffsetLeft, beginOffsetRight})

							anns = append(anns, sectionSegments[0]...)
							anns = append(anns, sectionSegments[1]...)
						}
					}
					lastDel, lastIns = -1, -1
				}
			}

			sort.Sort(anns)

			out, err := annotate.Annotate(src, anns, func(w io.Writer, b []byte) { template.HTMLEscape(w, b) })
			if err != nil {
				return nil, false
			}

			return out, true
		}
	default:
		return nil, false
	}
}

// Unexported blackfriday helpers.

func doubleSpace(out *bytes.Buffer) {
	if out.Len() > 0 {
		out.WriteByte('\n')
	}
}

func escapeSingleChar(char byte) (string, bool) {
	if char == '"' {
		return """, true
	}
	if char == '&' {
		return "&", true
	}
	if char == '<' {
		return "<", true
	}
	if char == '>' {
		return ">", true
	}
	return "", false
}

func attrEscape(out *bytes.Buffer, src []byte) {
	org := 0
	for i, ch := range src {
		if entity, ok := escapeSingleChar(ch); ok {
			if i > org {
				// copy all the normal characters since the last escape
				out.Write(src[org:i])
			}
			org = i + 1
			out.WriteString(entity)
		}
	}
	if org < len(src) {
		out.Write(src[org:])
	}
}
<commit_msg>Update moved import path.<commit_after>/*
Package github_flavored_markdown provides a GitHub Flavored Markdown renderer
with fenced code block highlighting, clickable header anchor links.

The functionality should be equivalent to the GitHub Markdown API endpoint specified at
https://developer.github.com/v3/markdown/#render-a-markdown-document-in-raw-mode, except
the rendering is performed locally.

See example below for how to generate a complete HTML page, including CSS styles.
*/
package github_flavored_markdown

import (
	"bytes"
	"fmt"
	"html"
	"io"
	"regexp"
	"sort"
	"strings"
	"text/template"

	"github.com/microcosm-cc/bluemonday"
	"github.com/russross/blackfriday"
	"github.com/shurcooL/sanitized_anchor_name"
	"github.com/shurcooL/go/highlight_go"
	"github.com/shurcooL/go/u/u7"
	"github.com/sourcegraph/annotate"
	"github.com/sourcegraph/syntaxhighlight"
)

// Markdown renders GitHub Flavored Markdown text.
func Markdown(text []byte) []byte {
	htmlFlags := 0
	renderer := &renderer{Html: blackfriday.HtmlRenderer(htmlFlags, "", "").(*blackfriday.Html)}

	// Parser extensions for GitHub Flavored Markdown.
	extensions := 0
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS
	//extensions |= blackfriday.EXTENSION_HARD_LINE_BREAK

	unsanitized := blackfriday.Markdown(text, renderer, extensions)

	// GitHub Flavored Markdown-like sanitization policy.
	p := bluemonday.UGCPolicy()
	p.AllowAttrs("class").Matching(bluemonday.SpaceSeparatedTokens).OnElements("div", "span")
	p.AllowAttrs("class", "name").Matching(bluemonday.SpaceSeparatedTokens).OnElements("a")
	p.AllowAttrs("rel").Matching(regexp.MustCompile(`^nofollow$`)).OnElements("a")
	p.AllowAttrs("aria-hidden").Matching(regexp.MustCompile(`^true$`)).OnElements("a")
	p.AllowDataURIImages()

	return p.SanitizeBytes(unsanitized)
}

type renderer struct {
	*blackfriday.Html
}

// GitHub Flavored Markdown header with clickable and hidden anchor.
func (_ *renderer) Header(out *bytes.Buffer, text func() bool, level int, _ string) {
	marker := out.Len()
	doubleSpace(out)

	if !text() {
		out.Truncate(marker)
		return
	}

	textString := out.String()[marker:]
	out.Truncate(marker)

	anchorName := sanitized_anchor_name.Create(html.UnescapeString(textString))

	out.WriteString(fmt.Sprintf(`<h%d><a name="%s" class="anchor" href="#%s" rel="nofollow" aria-hidden="true"><span class="octicon octicon-link"></span></a>`, level, anchorName, anchorName))
	out.WriteString(textString)
	out.WriteString(fmt.Sprintf("</h%d>\n", level))
}

// TODO: Clean up and improve this code.
// GitHub Flavored Markdown fenced code block with highlighting.
func (_ *renderer) BlockCode(out *bytes.Buffer, text []byte, lang string) {
	doubleSpace(out)

	// parse out the language name
	count := 0
	for _, elt := range strings.Fields(lang) {
		if elt[0] == '.' {
			elt = elt[1:]
		}
		if len(elt) == 0 {
			continue
		}
		out.WriteString(`<div class="highlight highlight-`)
		attrEscape(out, []byte(elt))
		lang = elt
		out.WriteString(`"><pre>`)
		count++
		break
	}

	if count == 0 {
		out.WriteString("
")
	}

	if highlightedCode, ok := highlightCode(text, lang); ok {
		out.Write(highlightedCode)
	} else {
		attrEscape(out, text)
	}

	if count == 0 {
		out.WriteString("
\n")
	} else {
		out.WriteString("

\n")
	}
}

var gfmHtmlConfig = syntaxhighlight.HTMLConfig{
	String:        "s",
	Keyword:       "k",
	Comment:       "c",
	Type:          "n",
	Literal:       "o",
	Punctuation:   "p",
	Plaintext:     "n",
	Tag:           "tag",
	HTMLTag:       "htm",
	HTMLAttrName:  "atn",
	HTMLAttrValue: "atv",
	Decimal:       "m",
}

// TODO: Support highlighting for more languages.
func highlightCode(src []byte, lang string) (highlightedCode []byte, ok bool) {
	switch lang {
	case "Go":
		var buf bytes.Buffer
		err := highlight_go.Print(src, &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "Go-old":
		var buf bytes.Buffer
		err := syntaxhighlight.Print(syntaxhighlight.NewScanner(src), &buf, syntaxhighlight.HTMLPrinter(gfmHtmlConfig))
		if err != nil {
			return nil, false
		}
		return buf.Bytes(), true
	case "diff":
		switch 2 {
		default:
			var buf bytes.Buffer
			err := u7.Print(u7.NewScanner(src), &buf)
			if err != nil {
				return nil, false
			}
			return buf.Bytes(), true
		case 1:
			lines := bytes.Split(src, []byte("\n"))
			return bytes.Join(lines, []byte("\n")), true
		case 2:
			anns, err := u7.Annotate(src)
			if err != nil {
				return nil, false
			}

			lines := bytes.Split(src, []byte("\n"))
			lineStarts := make([]int, len(lines))
			var offset int
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				lineStarts[lineIndex] = offset
				offset += len(lines[lineIndex]) + 1
			}

			lastDel, lastIns := -1, -1
			for lineIndex := 0; lineIndex < len(lines); lineIndex++ {
				var lineFirstChar byte
				if len(lines[lineIndex]) > 0 {
					lineFirstChar = lines[lineIndex][0]
				}
				switch lineFirstChar {
				case '+':
					if lastIns == -1 {
						lastIns = lineIndex
					}
				case '-':
					if lastDel == -1 {
						lastDel = lineIndex
					}
				default:
					if lastDel != -1 || lastIns != -1 {
						if lastDel == -1 {
							lastDel = lastIns
						} else if lastIns == -1 {
							lastIns = lineIndex
						}

						beginOffsetLeft := lineStarts[lastDel]
						endOffsetLeft := lineStarts[lastIns]
						beginOffsetRight := lineStarts[lastIns]
						endOffsetRight := lineStarts[lineIndex]

						anns = append(anns, &annotate.Annotation{Start: beginOffsetLeft, End: endOffsetLeft, Left: []byte(`<span class="gd input-block">`), Right: []byte(`</span>`), WantInner: 0})
						anns = append(anns, &annotate.Annotation{Start: beginOffsetRight, End: endOffsetRight, Left: []byte(`<span class="gi input-block">`), Right: []byte(`</span>`), WantInner: 0})

						if '@' != lineFirstChar {
							//leftContent := string(src[beginOffsetLeft:endOffsetLeft])
							//rightContent := string(src[beginOffsetRight:endOffsetRight])
							// This is needed to filter out the "-" and "+" at the beginning of each line from being highlighted.
							// TODO: Still not completely filtered out.
							leftContent := ""
							for line := lastDel; line < lastIns; line++ {
								leftContent += "\x00" + string(lines[line][1:]) + "\n"
							}
							rightContent := ""
							for line := lastIns; line < lineIndex; line++ {
								rightContent += "\x00" + string(lines[line][1:]) + "\n"
							}

							var sectionSegments [2][]*annotate.Annotation
							u7.HighlightedDiffFunc(leftContent, rightContent, &sectionSegments, [2]int{beginOffsetLeft, beginOffsetRight})

							anns = append(anns, sectionSegments[0]...)
							anns = append(anns, sectionSegments[1]...)
						}
					}
					lastDel, lastIns = -1, -1
				}
			}

			sort.Sort(anns)

			out, err := annotate.Annotate(src, anns, func(w io.Writer, b []byte) { template.HTMLEscape(w, b) })
			if err != nil {
				return nil, false
			}

			return out, true
		}
	default:
		return nil, false
	}
}

// Unexported blackfriday helpers.

func doubleSpace(out *bytes.Buffer) {
	if out.Len() > 0 {
		out.WriteByte('\n')
	}
}

func escapeSingleChar(char byte) (string, bool) {
	if char == '"' {
		return """, true
	}
	if char == '&' {
		return "&", true
	}
	if char == '<' {
		return "<", true
	}
	if char == '>' {
		return ">", true
	}
	return "", false
}

func attrEscape(out *bytes.Buffer, src []byte) {
	org := 0
	for i, ch := range src {
		if entity, ok := escapeSingleChar(ch); ok {
			if i > org {
				// copy all the normal characters since the last escape
				out.Write(src[org:i])
			}
			org = i + 1
			out.WriteString(entity)
		}
	}
	if org < len(src) {
		out.Write(src[org:])
	}
}
<|endoftext|>


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

package grpcproxy

import (
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

type ws2wc struct{ wserv pb.WatchServer }

func WatchServerToWatchClient(wserv pb.WatchServer) pb.WatchClient {
	return &ws2wc{wserv}
}

func (s *ws2wc) Watch(ctx context.Context, opts ...grpc.CallOption) (pb.Watch_WatchClient, error) {
	ch1, ch2 := make(chan interface{}), make(chan interface{})
	headerc, trailerc := make(chan metadata.MD, 1), make(chan metadata.MD, 1)
	wclient := &ws2wcClientStream{chanClientStream{headerc, trailerc, &chanStream{ch1, ch2, ctx}}}
	wserver := &ws2wcServerStream{chanServerStream{headerc, trailerc, &chanStream{ch2, ch1, ctx}}}
	go s.wserv.Watch(wserver)
	return wclient, nil
}

// ws2wcClientStream implements Watch_WatchClient
type ws2wcClientStream struct{ chanClientStream }

// ws2wcServerStream implements Watch_WatchServer
type ws2wcServerStream struct{ chanServerStream }

func (s *ws2wcClientStream) Send(wr *pb.WatchRequest) error {
	return s.SendMsg(wr)
}
func (s *ws2wcClientStream) Recv() (*pb.WatchResponse, error) {
	var v interface{}
	if err := s.RecvMsg(&v); err != nil {
		return nil, err
	}
	return v.(*pb.WatchResponse), nil
}

func (s *ws2wcServerStream) Send(wr *pb.WatchResponse) error {
	return s.SendMsg(wr)
}
func (s *ws2wcServerStream) Recv() (*pb.WatchRequest, error) {
	var v interface{}
	if err := s.RecvMsg(&v); err != nil {
		return nil, err
	}
	return v.(*pb.WatchRequest), nil
}

// chanServerStream implements grpc.ServerStream with a chanStream
type chanServerStream struct {
	headerc  chan<- metadata.MD
	trailerc chan<- metadata.MD
	grpc.Stream
}

func (ss *chanServerStream) SendHeader(md metadata.MD) error {
	select {
	case ss.headerc <- md:
		return nil
	case <-ss.Context().Done():
	}
	return ss.Context().Err()
}

func (ss *chanServerStream) SetTrailer(md metadata.MD) {
	ss.trailerc <- md
}

// chanClientStream implements grpc.ClientStream with a chanStream
type chanClientStream struct {
	headerc  <-chan metadata.MD
	trailerc <-chan metadata.MD
	grpc.Stream
}

func (cs *chanClientStream) Header() (metadata.MD, error) {
	select {
	case md := <-cs.headerc:
		return md, nil
	case <-cs.Context().Done():
	}
	return nil, cs.Context().Err()
}

func (cs *chanClientStream) Trailer() metadata.MD {
	select {
	case md := <-cs.trailerc:
		return md
	case <-cs.Context().Done():
		return nil
	}
}

func (s *chanClientStream) CloseSend() error { return nil }

// chanStream implements grpc.Stream using channels
type chanStream struct {
	recvc <-chan interface{}
	sendc chan<- interface{}
	ctx   context.Context
}

func (s *chanStream) Context() context.Context { return s.ctx }

func (s *chanStream) SendMsg(m interface{}) error {
	select {
	case s.sendc <- m:
		return nil
	case <-s.ctx.Done():
	}
	return s.ctx.Err()
}

func (s *chanStream) RecvMsg(m interface{}) error {
	v := m.(*interface{})
	select {
	case m = <-s.recvc:
		*v = m
		return nil
	case <-s.ctx.Done():
	}
	return s.ctx.Err()
}
<commit_msg>grpcproxy: fix stream closing issue<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grpcproxy

import (
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

type ws2wc struct{ wserv pb.WatchServer }

func WatchServerToWatchClient(wserv pb.WatchServer) pb.WatchClient {
	return &ws2wc{wserv}
}

func (s *ws2wc) Watch(ctx context.Context, opts ...grpc.CallOption) (pb.Watch_WatchClient, error) {
	ch1, ch2 := make(chan interface{}), make(chan interface{})
	headerc, trailerc := make(chan metadata.MD, 1), make(chan metadata.MD, 1)
	wclient := &ws2wcClientStream{chanClientStream{headerc, trailerc, &chanStream{ch1, ch2, ctx}}}
	wserver := &ws2wcServerStream{chanServerStream{headerc, trailerc, &chanStream{ch2, ch1, ctx}}}
	go func() {
		s.wserv.Watch(wserver)
		// close the server side sender
		close(ch1)
	}()
	return wclient, nil
}

// ws2wcClientStream implements Watch_WatchClient
type ws2wcClientStream struct{ chanClientStream }

// ws2wcServerStream implements Watch_WatchServer
type ws2wcServerStream struct{ chanServerStream }

func (s *ws2wcClientStream) Send(wr *pb.WatchRequest) error {
	return s.SendMsg(wr)
}
func (s *ws2wcClientStream) Recv() (*pb.WatchResponse, error) {
	var v interface{}
	if err := s.RecvMsg(&v); err != nil {
		return nil, err
	}
	return v.(*pb.WatchResponse), nil
}

func (s *ws2wcServerStream) Send(wr *pb.WatchResponse) error {
	return s.SendMsg(wr)
}
func (s *ws2wcServerStream) Recv() (*pb.WatchRequest, error) {
	var v interface{}
	if err := s.RecvMsg(&v); err != nil {
		return nil, err
	}
	return v.(*pb.WatchRequest), nil
}

// chanServerStream implements grpc.ServerStream with a chanStream
type chanServerStream struct {
	headerc  chan<- metadata.MD
	trailerc chan<- metadata.MD
	grpc.Stream
}

func (ss *chanServerStream) SendHeader(md metadata.MD) error {
	select {
	case ss.headerc <- md:
		return nil
	case <-ss.Context().Done():
	}
	return ss.Context().Err()
}

func (ss *chanServerStream) SetTrailer(md metadata.MD) {
	ss.trailerc <- md
}

// chanClientStream implements grpc.ClientStream with a chanStream
type chanClientStream struct {
	headerc  <-chan metadata.MD
	trailerc <-chan metadata.MD
	*chanStream
}

func (cs *chanClientStream) Header() (metadata.MD, error) {
	select {
	case md := <-cs.headerc:
		return md, nil
	case <-cs.Context().Done():
	}
	return nil, cs.Context().Err()
}

func (cs *chanClientStream) Trailer() metadata.MD {
	select {
	case md := <-cs.trailerc:
		return md
	case <-cs.Context().Done():
		return nil
	}
}

func (s *chanClientStream) CloseSend() error {
	close(s.chanStream.sendc)
	return nil
}

// chanStream implements grpc.Stream using channels
type chanStream struct {
	recvc <-chan interface{}
	sendc chan<- interface{}
	ctx   context.Context
}

func (s *chanStream) Context() context.Context { return s.ctx }

func (s *chanStream) SendMsg(m interface{}) error {
	select {
	case s.sendc <- m:
		return nil
	case <-s.ctx.Done():
	}
	return s.ctx.Err()
}

func (s *chanStream) RecvMsg(m interface{}) error {
	v := m.(*interface{})
	select {
	case msg, ok := <-s.recvc:
		if !ok {
			return grpc.ErrClientConnClosing
		}
		*v = msg
		return nil
	case <-s.ctx.Done():
	}
	return s.ctx.Err()
}
<|endoftext|>