<commit_before>package main

import (
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	exporterNamespace = "varnish"
)

// prometheusExporter

type prometheusExporter struct {
	sync.RWMutex

	up      prometheus.Gauge
	version prometheus.Gauge
}

func NewPrometheusExporter() *prometheusExporter {
	return &prometheusExporter{
		// @todo varnishstat never fails, even if varnish is not running.
		// Figure our a reliable way to detect if varnish is running and use that for up.
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: exporterNamespace,
			Name:      "up",
			Help:      "Was the last scrape of varnish successful.",
		}),
	}
}

func (pe *prometheusExporter) Initialize() error {
	pe.version = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace:   exporterNamespace,
		Name:        "version",
		Help:        "Varnish version information",
		ConstLabels: VarnishVersion.Labels(),
	})
	pe.version.Set(1)
	return nil
}

// Implements prometheus.Collector
func (pe *prometheusExporter) Describe(ch chan<- *prometheus.Desc) {
	start := time.Now()

	ch <- pe.up.Desc()
	ch <- pe.version.Desc()

	if StartParams.Verbose {
		logInfo("prometheus.Collector.Describe  %s", time.Now().Sub(start))
	}
}

// Implements prometheus.Collector
func (pe *prometheusExporter) Collect(ch chan<- prometheus.Metric) {
	start := time.Now()

	pe.Lock()
	defer pe.Unlock()

	if err := scrapeVarnish(ch); err == nil {
		pe.up.Set(1)
	} else {
		logError(err.Error())
		pe.up.Set(0)
	}

	ch <- pe.up
	ch <- pe.version

	if StartParams.Verbose {
		logInfo("prometheus.Collector.Collect   %s", time.Now().Sub(start))
	}
}

// utils

type group struct {
	name     string
	prefixes []string
}

var (
	groups = []group{
		// @note varnish 3.x does not seem to mark 'MAIN.' prefixes
		group{name: "backend", prefixes: []string{
			"vbe.",
			// varnish 4.x
			"main.backend_",
			// varnish 3.x
			"backend_",
		}},
		group{name: "mempool", prefixes: []string{
			"mempool.",
		}},
		group{name: "lck", prefixes: []string{
			"lck.",
		}},
		group{name: "sma", prefixes: []string{
			"sma.",
		}},
		group{name: "mgt", prefixes: []string{
			"mgt.",
		}},
		// must be last so above groups have a opportunity to override
		group{name: "main", prefixes: []string{
			"main.",
		}},
	}
)

type grouping struct {
	newPrefix string
	prefix    string
	total     string
	desc      string
}

var (
	fqGroupPrefixes = []*grouping{
		&grouping{
			prefix: "main_fetch",
			total:  "main_s_fetch",
			desc:   "Number of fetches",
		},
		&grouping{
			newPrefix: "main_sessions",
			prefix:    "main_sess",
			total:     "main_s_sess",
			desc:      "Number of sessions",
		},
		&grouping{
			newPrefix: "main_worker_threads",
			prefix:    "main_n_wrk",
			total:     "main_n_wrk",
			desc:      "Number of worker threads",
		},
		&grouping{
			newPrefix: "main_bans",
			prefix:    "main_n_ban",
			total:     "main_n_ban",
			desc:      "Number of bans",
		},
	}
)

// https://prometheus.io/docs/practices/naming/
func computePrometheusInfo(vName, vGroup, vIdentifier, vDescription string) (name, description string, labelKeys, labelValues []string) {
	// name and description
	{
		fq := strings.ToLower(vName)
		// Remove unique identifiers from name to group similar metrics by labeling
		if len(vIdentifier) > 0 {
			fq = strings.Replace(fq, "."+strings.ToLower(vIdentifier), "", -1)
		}
		// Make sure our group is prefixed only once
		fq = prometheusTrimGroupPrefix(fq)
		// Build fq name
		name = exporterNamespace + "_" + vGroup + "_" + strings.Replace(fq, ".", "_", -1)
		description = vDescription
	}
	// labels: can alter final name and description
	{
		if len(vIdentifier) > 0 {
			if isVBE := startsWith(vName, "VBE.", caseSensitive); isVBE {
				// @todo this is quick and dirty, do regexp?
				if VarnishVersion.Major == 4 {
					// <uuid>.<name>
					if len(vIdentifier) > 37 && vIdentifier[8] == '-' && vIdentifier[36] == '.' {
						labelKeys, labelValues = append(labelKeys, "server"), append(labelValues, vIdentifier[0:36])
						labelKeys, labelValues = append(labelKeys, "backend"), append(labelValues, vIdentifier[37:])
					}
				} else if VarnishVersion.Major == 3 {
					// <name>(<ip>,<something>,<port>)
					iStart, iEnd := strings.Index(vIdentifier, "("), strings.Index(vIdentifier, ")")
					if iStart > 0 && iEnd > 1 && iStart < iEnd {
						labelKeys, labelValues = append(labelKeys, "server"), append(labelValues, vIdentifier[iStart+1:iEnd])
						labelKeys, labelValues = append(labelKeys, "backend"), append(labelValues, vIdentifier[0:iStart])
					}
				}
			}
			if len(labelKeys) == 0 {
				labelKeys, labelValues = append(labelKeys, "id"), append(labelValues, vIdentifier)
			}
		}

		// create groupings by moving part of the fq name as a label and optional total
		for _, grouping := range fqGroupPrefixes {
			fqTotal := exporterNamespace + "_" + grouping.total
			fqPrefix := exporterNamespace + "_" + grouping.prefix
			fqNewName := fqPrefix
			if len(grouping.newPrefix) > 0 {
				fqNewName = exporterNamespace + "_" + grouping.newPrefix
			}
			if name == fqTotal {
				// @note total should not be a label as it breaks aggregation
				name, description = fqNewName+"_total", grouping.desc
				break
			} else if len(name) > len(fqPrefix)+1 && strings.HasPrefix(name, fqPrefix+"_") {
				labelKeys, labelValues = append(labelKeys, "type"), append(labelValues, name[len(fqPrefix)+1:])
				name, description = fqNewName, grouping.desc
				break
			}
		}
	}
	return name, description, labelKeys, labelValues
}

func prometheusTrimGroupPrefix(name string) string {
	nameLower := strings.ToLower(name)
	for _, group := range groups {
		for _, prefix := range group.prefixes {
			if startsWith(nameLower, prefix, caseSensitive) {
				return name[len(prefix):]
			}
		}
	}
	return name
}

// Always returns at least one main label
func prometheusGroup(vName string) string {
	vNameLower := strings.ToLower(vName)
	for _, group := range groups {
		if startsWithAny(vNameLower, group.prefixes, caseSensitive) {
			return group.name
		}
	}
	return "main"
}
<commit_msg>Implement renaming of names and label keys. Still some generic 'type' ones but no more 'id'.<commit_after>package main

import (
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	exporterNamespace = "varnish"
)

// prometheusExporter

type prometheusExporter struct {
	sync.RWMutex

	up      prometheus.Gauge
	version prometheus.Gauge
}

func NewPrometheusExporter() *prometheusExporter {
	return &prometheusExporter{
		// @todo varnishstat never fails, even if varnish is not running.
		// Figure our a reliable way to detect if varnish is running and use that for up.
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: exporterNamespace,
			Name:      "up",
			Help:      "Was the last scrape of varnish successful.",
		}),
	}
}

func (pe *prometheusExporter) Initialize() error {
	pe.version = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace:   exporterNamespace,
		Name:        "version",
		Help:        "Varnish version information",
		ConstLabels: VarnishVersion.Labels(),
	})
	pe.version.Set(1)
	return nil
}

// Implements prometheus.Collector
func (pe *prometheusExporter) Describe(ch chan<- *prometheus.Desc) {
	start := time.Now()

	ch <- pe.up.Desc()
	ch <- pe.version.Desc()

	if StartParams.Verbose {
		logInfo("prometheus.Collector.Describe  %s", time.Now().Sub(start))
	}
}

// Implements prometheus.Collector
func (pe *prometheusExporter) Collect(ch chan<- prometheus.Metric) {
	start := time.Now()

	pe.Lock()
	defer pe.Unlock()

	if err := scrapeVarnish(ch); err == nil {
		pe.up.Set(1)
	} else {
		logError(err.Error())
		pe.up.Set(0)
	}

	ch <- pe.up
	ch <- pe.version

	if StartParams.Verbose {
		logInfo("prometheus.Collector.Collect   %s", time.Now().Sub(start))
	}
}

// utils

type group struct {
	name     string
	prefixes []string
}

var (
	groups = []group{
		// @note varnish 3.x does not seem to mark 'MAIN.' prefixes
		group{name: "backend", prefixes: []string{
			"vbe.",
			// varnish 4.x
			"main.backend_",
			// varnish 3.x
			"backend_",
		}},
		group{name: "mempool", prefixes: []string{
			"mempool.",
		}},
		group{name: "lck", prefixes: []string{
			"lck.",
		}},
		group{name: "sma", prefixes: []string{
			"sma.",
		}},
		group{name: "mgt", prefixes: []string{
			"mgt.",
		}},
		// must be last so above groups have a opportunity to override
		group{name: "main", prefixes: []string{
			"main.",
		}},
	}
)

type grouping struct {
	newPrefix string
	prefix    string
	total     string
	desc      string
	labelKey  string
}

var (
	fqGroupPrefixes = []*grouping{
		&grouping{
			prefix: "main_fetch",
			total:  "main_s_fetch",
			desc:   "Number of fetches",
		},
		&grouping{
			newPrefix: "main_sessions",
			prefix:    "main_sess",
			total:     "main_s_sess",
			desc:      "Number of sessions",
		},
		&grouping{
			newPrefix: "main_worker_threads",
			prefix:    "main_n_wrk",
			total:     "main_n_wrk",
			desc:      "Number of worker threads",
		},
		&grouping{
			newPrefix: "main_bans",
			prefix:    "main_n_ban",
			total:     "main_n_ban",
			desc:      "Number of bans",
			labelKey:  "action",
		},
	}
	fqNames = map[string]string{
		"varnish_lck_colls":   "varnish_lock_collisions",
		"varnish_lck_creat":   "varnish_lock_created",
		"varnish_lck_destroy": "varnish_lock_destroyed",
		"varnish_lck_locks":   "varnish_lock_operations",
	}
	fqIdentifiers = map[string]string{
		"varnish_lock_collisions": "target",
		"varnish_lock_created":    "target",
		"varnish_lock_destroyed":  "target",
		"varnish_lock_operations": "target",
		"varnish_sma_c_bytes":     "type",
		"varnish_sma_c_fail":      "type",
		"varnish_sma_c_freed":     "type",
		"varnish_sma_c_req":       "type",
		"varnish_sma_g_alloc":     "type",
		"varnish_sma_g_bytes":     "type",
		"varnish_sma_g_space":     "type",
	}
)

// https://prometheus.io/docs/practices/naming/
func computePrometheusInfo(vName, vGroup, vIdentifier, vDescription string) (name, description string, labelKeys, labelValues []string) {
	// name and description
	{
		fq := strings.ToLower(vName)
		// Remove unique identifiers from name to group similar metrics by labeling
		if len(vIdentifier) > 0 {
			fq = strings.Replace(fq, "."+strings.ToLower(vIdentifier), "", -1)
		}
		// Make sure our group is prefixed only once
		fq = prometheusTrimGroupPrefix(fq)
		// Build fq name
		name = exporterNamespace + "_" + vGroup + "_" + strings.Replace(fq, ".", "_", -1)
		if swapName := fqNames[name]; len(swapName) > 0 {
			name = swapName
		}
		description = vDescription
	}
	// labels: can alter final name and description
	{
		if len(vIdentifier) > 0 {
			if isVBE := startsWith(vName, "VBE.", caseSensitive); isVBE {
				// @todo this is quick and dirty, do regexp?
				if VarnishVersion.Major == 4 {
					// <uuid>.<name>
					if len(vIdentifier) > 37 && vIdentifier[8] == '-' && vIdentifier[36] == '.' {
						labelKeys, labelValues = append(labelKeys, "server"), append(labelValues, vIdentifier[0:36])
						labelKeys, labelValues = append(labelKeys, "backend"), append(labelValues, vIdentifier[37:])
					}
				} else if VarnishVersion.Major == 3 {
					// <name>(<ip>,<something>,<port>)
					iStart, iEnd := strings.Index(vIdentifier, "("), strings.Index(vIdentifier, ")")
					if iStart > 0 && iEnd > 1 && iStart < iEnd {
						labelKeys, labelValues = append(labelKeys, "server"), append(labelValues, vIdentifier[iStart+1:iEnd])
						labelKeys, labelValues = append(labelKeys, "backend"), append(labelValues, vIdentifier[0:iStart])
					}
				}
			}
			if len(labelKeys) == 0 {
				labelKey := fqIdentifiers[name]
				if len(labelKey) == 0 {
					labelKey = "id"
				}
				labelKeys, labelValues = append(labelKeys, labelKey), append(labelValues, vIdentifier)
			}
		}

		// create groupings by moving part of the fq name as a label and optional total
		for _, grouping := range fqGroupPrefixes {
			fqTotal := exporterNamespace + "_" + grouping.total
			fqPrefix := exporterNamespace + "_" + grouping.prefix
			fqNewName := fqPrefix
			if len(grouping.newPrefix) > 0 {
				fqNewName = exporterNamespace + "_" + grouping.newPrefix
			}
			if name == fqTotal {
				// @note total should not be a label as it breaks aggregation
				name, description = fqNewName+"_total", grouping.desc
				break
			} else if len(name) > len(fqPrefix)+1 && strings.HasPrefix(name, fqPrefix+"_") {
				labelKey := "type"
				if len(grouping.labelKey) > 0 {
					labelKey = grouping.labelKey
				}
				labelKeys, labelValues = append(labelKeys, labelKey), append(labelValues, name[len(fqPrefix)+1:])
				name, description = fqNewName, grouping.desc
				break
			}
		}
	}
	return name, description, labelKeys, labelValues
}

func prometheusTrimGroupPrefix(name string) string {
	nameLower := strings.ToLower(name)
	for _, group := range groups {
		for _, prefix := range group.prefixes {
			if startsWith(nameLower, prefix, caseSensitive) {
				return name[len(prefix):]
			}
		}
	}
	return name
}

// Always returns at least one main label
func prometheusGroup(vName string) string {
	vNameLower := strings.ToLower(vName)
	for _, group := range groups {
		if startsWithAny(vNameLower, group.prefixes, caseSensitive) {
			return group.name
		}
	}
	return "main"
}
<|endoftext|>


<commit_before>package CSVWriter

import (
	. "github.com/xitongsys/parquet-go/Common"
	. "github.com/xitongsys/parquet-go/SchemaHandler"
	"github.com/xitongsys/parquet-go/parquet"
)

func MarshalCSV(records [][]*string, bgn int, end int, md []MetadataType, schemaHandler *SchemaHandler) *map[string]*Table {
	res := make(map[string]*Table)
	for i := 0; i < len(md); i++ {
		pathStr := "parquet-go-root." + md[i].Name
		res[pathStr] = new(Table)
		res[pathStr].Path = StrToPath(pathStr)
		res[pathStr].MaxDefinitionLevel = 1
		res[pathStr].MaxRepetitionLevel = 0
		res[pathStr].Repetition_Type = parquet.FieldRepetitionType_OPTIONAL
		res[pathStr].Type = schemaHandler.SchemaElements[schemaHandler.MapIndex[pathStr]].GetType()

		for j := bgn; j < end; j++ {
			rec := records[j][i]
			res[pathStr].Values = append(res[pathStr].Values, rec)
			res[pathStr].RepetitionLevels = append(res[pathStr].RepetitionLevels, 0)

			if rec == nil {
				res[pathStr].DefinitionLevels = append(res[pathStr].DefinitionLevels, 0)
			} else {
				res[pathStr].DefinitionLevels = append(res[pathStr].DefinitionLevels, 1)
			}
		}
	}
	return &res
}
<commit_msg>adding csvwriter: fix marshcsv bug<commit_after>package CSVWriter

import (
	. "github.com/xitongsys/parquet-go/Common"
	. "github.com/xitongsys/parquet-go/ParquetType"
	. "github.com/xitongsys/parquet-go/SchemaHandler"
	"github.com/xitongsys/parquet-go/parquet"
)

func MarshalCSV(records [][]*string, bgn int, end int, md []MetadataType, schemaHandler *SchemaHandler) *map[string]*Table {
	res := make(map[string]*Table)
	for i := 0; i < len(md); i++ {
		pathStr := "parquet-go-root." + md[i].Name
		res[pathStr] = new(Table)
		res[pathStr].Path = StrToPath(pathStr)
		res[pathStr].MaxDefinitionLevel = 1
		res[pathStr].MaxRepetitionLevel = 0
		res[pathStr].Repetition_Type = parquet.FieldRepetitionType_OPTIONAL
		res[pathStr].Type = schemaHandler.SchemaElements[schemaHandler.MapIndex[pathStr]].GetType()

		for j := bgn; j < end; j++ {
			rec := StrToParquetType(*records[j][i], md[i].Type)
			res[pathStr].Values = append(res[pathStr].Values, rec)
			res[pathStr].RepetitionLevels = append(res[pathStr].RepetitionLevels, 0)

			if rec == nil {
				res[pathStr].DefinitionLevels = append(res[pathStr].DefinitionLevels, 0)
			} else {
				res[pathStr].DefinitionLevels = append(res[pathStr].DefinitionLevels, 1)
			}
		}
	}
	return &res
}
<|endoftext|>


<commit_before>package helpers

import (
	. "github.com/onsi/gomega"
	. "github.com/vito/cmdtest/matchers"
	"time"
)

func PushApp(appName string, appPath string) {
	Expect(Cf("push", appName, "-p", appPath)).To(SayWithTimeout("App started", time.Minute*1))
}

func DeleteApp(appName string) {
	Expect(Cf("delete", appName, "-f")).To(SayWithTimeout("OK", time.Minute*1))
}
<commit_msg>increase timeout for pushing an app<commit_after>package helpers

import (
	. "github.com/onsi/gomega"
	. "github.com/vito/cmdtest/matchers"
	"time"
)

func PushApp(appName string, appPath string) {
	Expect(Cf("push", appName, "-p", appPath)).To(SayWithTimeout("App started", time.Minute*2))
}

func DeleteApp(appName string) {
	Expect(Cf("delete", appName, "-f")).To(SayWithTimeout("OK", time.Minute*2))
}
<|endoftext|>


<commit_before>// Copyright 2014 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package boulder

import (
	"crypto/sha256"
	"crypto/x509"
	"fmt"
	"github.com/bifurcation/gose"
	"regexp"
	"strings"
)

// All of the fields in RegistrationAuthorityImpl need to be
// populated, or there is a risk of panic.
type RegistrationAuthorityImpl struct {
	CA CertificateAuthority
	VA ValidationAuthority
	SA StorageAuthority
}

func NewRegistrationAuthorityImpl() RegistrationAuthorityImpl {
	return RegistrationAuthorityImpl{}
}

var dnsLabelRegexp, _ = regexp.Compile("[^a-zA-Z0-9-]")
var ipAddressRegexp, _ = regexp.Compile("^[0-9.]*$")

func forbiddenIdentifier(id string) bool {
	// A DNS label is a part separated by dots, e.g. www.foo.net has labels
	// "www", "foo", and "net".
	const maxLabels = 10
	labels := strings.SplitN(id, ".", maxLabels + 1)
	if len(labels) < 2 || len(labels) > maxLabels {
		return true
	}

	for _, label := range labels {
		// DNS defines max label length as 63 characters. Some implementations allow
		// more, but we will be conservative.
		if len(label) < 1 || len(label) > 63 {
			return true
		}
		// Only alphanumerics and dash are allowed in identifiers.
		// TODO: Before identifiers reach this function, do lowercasing.
		if dnsLabelRegexp.MatchString(label) {
			return true
		}

		// A label cannot begin with a hyphen (-)
		if label[0] == '-' {
			return true
		}

		// Punycode labels are not yet allowed. May allow in future after looking at
		// homoglyph mitigations.
		if len(label) >= 4 && label[0:4] == "xn--" {
			return true
		}
	}

	// Forbid identifiers that are entirely numeric like an IP address.
	if ipAddressRegexp.MatchString(id) {
		return true
	}

	// Also forbid an all-numeric final label.
	if ipAddressRegexp.MatchString(labels[len(labels) - 1]) {
		return true
	}

	return false
}

func fingerprint256(data []byte) string {
	d := sha256.New()
	d.Write(data)
	return b64enc(d.Sum(nil))
}

func (ra *RegistrationAuthorityImpl) NewAuthorization(request Authorization, key jose.JsonWebKey) (Authorization, error) {
	zero := Authorization{}
	identifier := request.Identifier

	// Check that the identifier is present and appropriate
	if len(identifier.Value) == 0 {
		return zero, MalformedRequestError("No identifier in authorization request")
	} else if identifier.Type != IdentifierDNS {
		return zero, NotSupportedError("Only domain validation is supported")
	} else if forbiddenIdentifier(identifier.Value) {
		return zero, UnauthorizedError("We will not authorize use of this identifier")
	}

	// Create validations
	authID := newToken()
	simpleHttps := SimpleHTTPSChallenge()
	dvsni := DvsniChallenge()

	// Create a new authorization object
	authz := Authorization{
		ID:         authID,
		Identifier: identifier,
		Key:        key,
		Status:     StatusPending,
		Challenges: map[string]Challenge{
			ChallengeTypeSimpleHTTPS: simpleHttps,
			ChallengeTypeDVSNI:       dvsni,
		},
	}

	// Store the authorization object, then return it
	err := ra.SA.Update(authz.ID, authz)
	if err != nil {
		return authz, err
	}

	return authz, nil
}

func (ra *RegistrationAuthorityImpl) NewCertificate(req CertificateRequest, jwk jose.JsonWebKey) (Certificate, error) {
	csr := req.CSR
	zero := Certificate{}

	// Verify the CSR
	// TODO: Verify that other aspects of the CSR are appropriate
	err := VerifyCSR(csr)
	if err != nil {
		return zero, UnauthorizedError("Invalid signature on CSR")
	}

	// Get the authorized domain list for the authorization key
	obj, err := ra.SA.Get(jwk.Thumbprint)
	if err != nil {
		return zero, UnauthorizedError("No authorized domains for this key")
	}
	domainSet := obj.(map[string]bool)

	// Validate that authorization key is authorized for all domains
	// TODO: Use linked authorizations?
	names := csr.DNSNames
	if len(csr.Subject.CommonName) > 0 {
		names = append(names, csr.Subject.CommonName)
	}
	for _, name := range names {
		if !domainSet[name] {
			return zero, UnauthorizedError(fmt.Sprintf("Key not authorized for name %s", name))
		}
	}

	// Create the certificate
	cert, err := ra.CA.IssueCertificate(*csr)
	if err != nil {
		return zero, CertificateIssuanceError("Error issuing certificate")
	}

	// Identify the certificate object by the cert's SHA-256 fingerprint
	certObj := Certificate{
		ID:     fingerprint256(cert),
		DER:    cert,
		Status: StatusValid,
	}

	ra.SA.Update(certObj.ID, certObj)
	return certObj, nil
}

func (ra *RegistrationAuthorityImpl) UpdateAuthorization(delta Authorization) (Authorization, error) {
	zero := Authorization{}

	// Fetch the copy of this authorization we have on file
	obj, err := ra.SA.Get(delta.ID)
	if err != nil {
		return zero, err
	}
	authz := obj.(Authorization)

	// Copy information over that the client is allowed to supply
	if len(delta.Contact) > 0 {
		authz.Contact = delta.Contact
	}
	newResponse := false
	for t, challenge := range authz.Challenges {
		response, present := delta.Challenges[t]
		if !present {
			continue
		}

		newResponse = true
		authz.Challenges[t] = challenge.MergeResponse(response)
	}

	// Store the updated version
	ra.SA.Update(authz.ID, authz)

	// If any challenges were updated, dispatch to the VA for service
	if newResponse {
		ra.VA.UpdateValidations(authz)
	}

	return authz, nil
}

func (ra *RegistrationAuthorityImpl) RevokeCertificate(cert x509.Certificate) error {
	// Attempt to fetch the corresponding certificate object
	certID := fingerprint256(cert.Raw)
	obj, err := ra.SA.Get(certID)
	if err != nil {
		return err
	}
	certObj := obj.(Certificate)

	// Change the status and update the DB
	certObj.Status = StatusInvalid
	return ra.SA.Update(certID, certObj)
}

func (ra *RegistrationAuthorityImpl) OnValidationUpdate(authz Authorization) {
	// Check to see whether the updated validations are sufficient
	// Current policy is to accept if any validation succeeded
	for _, val := range authz.Challenges {
		if val.Status == StatusValid {
			authz.Status = StatusValid
			break
		}
	}

	// If no validation succeeded, then the authorization is invalid
	// NOTE: This only works because we only ever do one validation
	if authz.Status != StatusValid {
		authz.Status = StatusInvalid
	}
	ra.SA.Update(authz.ID, authz)

	// Record a new domain/key binding, if authorized
	if authz.Status == StatusValid {
		var domainSet map[string]bool
		obj, err := ra.SA.Get(authz.Key.Thumbprint)
		if err != nil {
			domainSet = make(map[string]bool)
		} else {
			domainSet = obj.(map[string]bool)
		}
		domainSet[authz.Identifier.Value] = true
		ra.SA.Update(authz.Key.Thumbprint, domainSet)
	}
}
<commit_msg>Invert sense of dnsLabelRegexp.<commit_after>// Copyright 2014 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package boulder

import (
	"crypto/sha256"
	"crypto/x509"
	"fmt"
	"github.com/bifurcation/gose"
	"regexp"
	"strings"
)

// All of the fields in RegistrationAuthorityImpl need to be
// populated, or there is a risk of panic.
type RegistrationAuthorityImpl struct {
	CA CertificateAuthority
	VA ValidationAuthority
	SA StorageAuthority
}

func NewRegistrationAuthorityImpl() RegistrationAuthorityImpl {
	return RegistrationAuthorityImpl{}
}

var dnsLabelRegexp, _ = regexp.Compile("^[a-zA-Z0-9-]*$")
var ipAddressRegexp, _ = regexp.Compile("^[0-9.]*$")

func forbiddenIdentifier(id string) bool {
	// A DNS label is a part separated by dots, e.g. www.foo.net has labels
	// "www", "foo", and "net".
	const maxLabels = 10
	labels := strings.SplitN(id, ".", maxLabels + 1)
	if len(labels) < 2 || len(labels) > maxLabels {
		return true
	}

	for _, label := range labels {
		// DNS defines max label length as 63 characters. Some implementations allow
		// more, but we will be conservative.
		if len(label) < 1 || len(label) > 63 {
			return true
		}
		// Only alphanumerics and dash are allowed in identifiers.
		// TODO: Before identifiers reach this function, do lowercasing.
		if ! dnsLabelRegexp.MatchString(label) {
			return true
		}

		// A label cannot begin with a hyphen (-)
		if label[0] == '-' {
			return true
		}

		// Punycode labels are not yet allowed. May allow in future after looking at
		// homoglyph mitigations.
		if len(label) >= 4 && label[0:4] == "xn--" {
			return true
		}
	}

	// Forbid identifiers that are entirely numeric like an IP address.
	if ipAddressRegexp.MatchString(id) {
		return true
	}

	// Also forbid an all-numeric final label.
	if ipAddressRegexp.MatchString(labels[len(labels) - 1]) {
		return true
	}

	return false
}

func fingerprint256(data []byte) string {
	d := sha256.New()
	d.Write(data)
	return b64enc(d.Sum(nil))
}

func (ra *RegistrationAuthorityImpl) NewAuthorization(request Authorization, key jose.JsonWebKey) (Authorization, error) {
	zero := Authorization{}
	identifier := request.Identifier

	// Check that the identifier is present and appropriate
	if len(identifier.Value) == 0 {
		return zero, MalformedRequestError("No identifier in authorization request")
	} else if identifier.Type != IdentifierDNS {
		return zero, NotSupportedError("Only domain validation is supported")
	} else if forbiddenIdentifier(identifier.Value) {
		return zero, UnauthorizedError("We will not authorize use of this identifier")
	}

	// Create validations
	authID := newToken()
	simpleHttps := SimpleHTTPSChallenge()
	dvsni := DvsniChallenge()

	// Create a new authorization object
	authz := Authorization{
		ID:         authID,
		Identifier: identifier,
		Key:        key,
		Status:     StatusPending,
		Challenges: map[string]Challenge{
			ChallengeTypeSimpleHTTPS: simpleHttps,
			ChallengeTypeDVSNI:       dvsni,
		},
	}

	// Store the authorization object, then return it
	err := ra.SA.Update(authz.ID, authz)
	if err != nil {
		return authz, err
	}

	return authz, nil
}

func (ra *RegistrationAuthorityImpl) NewCertificate(req CertificateRequest, jwk jose.JsonWebKey) (Certificate, error) {
	csr := req.CSR
	zero := Certificate{}

	// Verify the CSR
	// TODO: Verify that other aspects of the CSR are appropriate
	err := VerifyCSR(csr)
	if err != nil {
		return zero, UnauthorizedError("Invalid signature on CSR")
	}

	// Get the authorized domain list for the authorization key
	obj, err := ra.SA.Get(jwk.Thumbprint)
	if err != nil {
		return zero, UnauthorizedError("No authorized domains for this key")
	}
	domainSet := obj.(map[string]bool)

	// Validate that authorization key is authorized for all domains
	// TODO: Use linked authorizations?
	names := csr.DNSNames
	if len(csr.Subject.CommonName) > 0 {
		names = append(names, csr.Subject.CommonName)
	}
	for _, name := range names {
		if !domainSet[name] {
			return zero, UnauthorizedError(fmt.Sprintf("Key not authorized for name %s", name))
		}
	}

	// Create the certificate
	cert, err := ra.CA.IssueCertificate(*csr)
	if err != nil {
		return zero, CertificateIssuanceError("Error issuing certificate")
	}

	// Identify the certificate object by the cert's SHA-256 fingerprint
	certObj := Certificate{
		ID:     fingerprint256(cert),
		DER:    cert,
		Status: StatusValid,
	}

	ra.SA.Update(certObj.ID, certObj)
	return certObj, nil
}

func (ra *RegistrationAuthorityImpl) UpdateAuthorization(delta Authorization) (Authorization, error) {
	zero := Authorization{}

	// Fetch the copy of this authorization we have on file
	obj, err := ra.SA.Get(delta.ID)
	if err != nil {
		return zero, err
	}
	authz := obj.(Authorization)

	// Copy information over that the client is allowed to supply
	if len(delta.Contact) > 0 {
		authz.Contact = delta.Contact
	}
	newResponse := false
	for t, challenge := range authz.Challenges {
		response, present := delta.Challenges[t]
		if !present {
			continue
		}

		newResponse = true
		authz.Challenges[t] = challenge.MergeResponse(response)
	}

	// Store the updated version
	ra.SA.Update(authz.ID, authz)

	// If any challenges were updated, dispatch to the VA for service
	if newResponse {
		ra.VA.UpdateValidations(authz)
	}

	return authz, nil
}

func (ra *RegistrationAuthorityImpl) RevokeCertificate(cert x509.Certificate) error {
	// Attempt to fetch the corresponding certificate object
	certID := fingerprint256(cert.Raw)
	obj, err := ra.SA.Get(certID)
	if err != nil {
		return err
	}
	certObj := obj.(Certificate)

	// Change the status and update the DB
	certObj.Status = StatusInvalid
	return ra.SA.Update(certID, certObj)
}

func (ra *RegistrationAuthorityImpl) OnValidationUpdate(authz Authorization) {
	// Check to see whether the updated validations are sufficient
	// Current policy is to accept if any validation succeeded
	for _, val := range authz.Challenges {
		if val.Status == StatusValid {
			authz.Status = StatusValid
			break
		}
	}

	// If no validation succeeded, then the authorization is invalid
	// NOTE: This only works because we only ever do one validation
	if authz.Status != StatusValid {
		authz.Status = StatusInvalid
	}
	ra.SA.Update(authz.ID, authz)

	// Record a new domain/key binding, if authorized
	if authz.Status == StatusValid {
		var domainSet map[string]bool
		obj, err := ra.SA.Get(authz.Key.Thumbprint)
		if err != nil {
			domainSet = make(map[string]bool)
		} else {
			domainSet = obj.(map[string]bool)
		}
		domainSet[authz.Identifier.Value] = true
		ra.SA.Update(authz.Key.Thumbprint, domainSet)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 beego 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.
//
// Maintain by https://github.com/slene

package social

import (
	"encoding/base64"
	"fmt"
	"net/url"
	"strconv"

	"github.com/astaxie/beego"
	"github.com/astaxie/beego/context"
	"github.com/astaxie/beego/utils"
)

const (
	defaultURLPrefix          = "/loginsocial/"
	defaultConnectSuccessURL  = "/?flag=connect_success"
	defaultConnectFailedURL   = "/login?flag=connect_failed"
	defaultLoginURL           = "/loginsocial"
	defaultConnectRegisterURL = "/register/connect"
)

type SocialAuth struct {
	app                SocialAuther
	URLPrefix          string
	ConnectSuccessURL  string
	ConnectFailedURL   string
	LoginURL           string
	ConnectRegisterURL string
}

// generate session key for social-auth
func (this *SocialAuth) getSessKey(social SocialType, key string) string {
	return "social_" + fmt.Sprintf("%v", social) + "_" + key
}

// create oauth2 state string
func (this *SocialAuth) createState(ctx *context.Context, social SocialType) string {
	values := make(url.Values, 2)

	if uid, ok := this.app.IsUserLogin(ctx); ok {
		// add uid if user current is login
		values.Add("uid", strconv.FormatInt(int64(uid), 10))
	}

	// our secret string
	values.Add("secret", string(utils.RandomCreateBytes(20)))

	// create state string
	state := base64.URLEncoding.EncodeToString([]byte(values.Encode()))

	// save to session
	name := this.getSessKey(social, "state")
	ctx.Input.CruSession.Set(name, state)

	return state
}

// verify oauth2 state string
func (this *SocialAuth) verifyState(ctx *context.Context, social SocialType) (string, bool) {
	code := ctx.Input.Query("code")
	state := ctx.Input.Query("state")

	if len(code) == 0 || len(state) == 0 {
		return "", false
	}

	name := this.getSessKey(social, "state")

	vu, ok := ctx.Input.CruSession.Get(name).(string)
	if !ok || ok && state != vu {
		return "", false
	}

	return code, true
}

// Get provider according request path. ex: /login/: match /login/github
func (this *SocialAuth) getProvider(ctx *context.Context) Provider {
	path := ctx.Input.Param(":splat")

	p, ok := GetProviderByPath(path)
	if ok {
		return p
	}

	return nil
}

// After OAuthAccess check saved token for ready connect
func (this *SocialAuth) ReadyConnect(ctx *context.Context) (SocialType, bool) {
	var social SocialType

	if s, _ := ctx.Input.CruSession.Get("social_connect").(int); s == 0 {
		return 0, false
	} else {
		social = SocialType(s)
	}

	if !social.Available() {
		return 0, false
	}

	return social, true
}

// Redirect to other social platform
func (this *SocialAuth) OAuthRedirect(ctx *context.Context) (redirect string, failedErr error) {
	_, isLogin := this.app.IsUserLogin(ctx)

	defer func() {
		if len(redirect) == 0 && failedErr != nil {
			if isLogin {
				redirect = this.ConnectFailedURL
			} else {
				redirect = this.LoginURL
			}
		}
	}()

	var p Provider
	if p = this.getProvider(ctx); p == nil {
		failedErr = fmt.Errorf("unknown provider")
		return
	}

	social := p.GetType()
	config := p.GetConfig()
	// create redirect url
	redirect = config.AuthCodeURL(this.createState(ctx, social))
	return
}

// Callback from social platform
func (this *SocialAuth) OAuthAccess(ctx *context.Context) (redirect string, userSocial *UserSocial, failedErr error) {
	_, isLogin := this.app.IsUserLogin(ctx)

	defer func() {
		if len(redirect) == 0 {
			if failedErr != nil {
				if isLogin {
					redirect = this.ConnectFailedURL
				} else {
					redirect = this.LoginURL
				}
			}
		}
	}()

	// check if param has a error key
	if err := ctx.Input.Query("error"); len(err) > 0 {
		failedErr = fmt.Errorf(err)
		return
	}

	// get provider from matched url path
	var p Provider
	if p = this.getProvider(ctx); p == nil {
		failedErr = fmt.Errorf("unknown provider")
		return
	}

	social := p.GetType()

	var code string

	// verify state string
	if c, ok := this.verifyState(ctx, social); !ok {
		failedErr = fmt.Errorf("state not verified")
		return
	} else {
		code = c
	}

	config := p.GetConfig()
	trans := &Transport{config, nil, DefaultTransport}

	// Send code to platform then get token
	if tok, err := trans.Exchange(code); err != nil {
		// get access token
		failedErr = err
	} else if err := tok.GetExtra("error"); err != "" {
		// token has error
		failedErr = fmt.Errorf(err)
	} else if tok.IsEmpty() {
		failedErr = fmt.Errorf("empty access token")
	} else {

		// check
		var uSocial = UserSocial{}
		if ok, err := p.CanConnect(tok, &uSocial); ok {
			// save token to session, for connect
			tk := SocialTokenField{tok}
			ctx.Input.CruSession.Set(this.getSessKey(social, "token"), tk.RawValue())
			ctx.Input.CruSession.Set("social_connect", int(social))

			redirect = this.ConnectRegisterURL

		} else if err == nil {
			if !isLogin {
				// login user
				redirect, failedErr = this.app.LoginUser(ctx, uSocial.Uid)
				if len(redirect) == 0 && failedErr == nil {
					redirect = this.ConnectSuccessURL
				}
			} else {
				redirect = this.ConnectSuccessURL
			}

			// save new access token if it changed
			uSocial.PutToken(tok)

			userSocial = &uSocial

		} else {
			failedErr = err
		}
	}

	return
}

// general use of redirect
func (this *SocialAuth) handleRedirect(ctx *context.Context) {
	redirect, err := this.OAuthRedirect(ctx)
	if err != nil {
		beego.Error("SocialAuth.handleRedirect", err)
	}

	if len(redirect) > 0 {
		ctx.Redirect(302, redirect)
	}
}

// general use of redirect callback
func (this *SocialAuth) handleAccess(ctx *context.Context) {
	redirect, _, err := this.OAuthAccess(ctx)
	if err != nil {
		beego.Error("SocialAuth.handleAccess", err)
	}

	if len(redirect) > 0 {
		ctx.Redirect(302, redirect)
	}
}

// save user social info and login the user
func (this *SocialAuth) ConnectAndLogin(ctx *context.Context, socialType SocialType, uid int64) (string, *UserSocial, error) {
	tokKey := this.getSessKey(socialType, "token")

	defer func() {
		// delete connect tok in session
		if ctx.Input.CruSession.Get("social_connect") != nil {
			ctx.Input.CruSession.Delete("social_connect")
		}
		if ctx.Input.CruSession.Get(tokKey) != nil {
			ctx.Input.CruSession.Delete(tokKey)
		}
	}()

	tk := SocialTokenField{}
	value := ctx.Input.CruSession.Get(tokKey)
	if err := tk.SetRaw(value); err != nil {
		return "", nil, err
	}

	var p Provider
	if p, _ = GetProviderByType(socialType); p == nil {
		return "", nil, fmt.Errorf("unknown provider")
	}

	identify, err := p.GetIndentify(tk.Token)
	if err != nil {
		return "", nil, err
	}
	if len(identify) == 0 {
		return "", nil, fmt.Errorf("empty identify")
	}

	userSocial := UserSocial{
		Uid:      uid,
		Type:     socialType,
		Data:     tk,
		Identify: identify,
	}

	if err := userSocial.Save(); err != nil {
		return "", nil, err
	}

	// login user
	loginRedirect, err := this.app.LoginUser(ctx, uid)
	return loginRedirect, &userSocial, nil
}

// create a global SocialAuth instance
func NewSocial(urlPrefix string, socialAuther SocialAuther) *SocialAuth {
	social := new(SocialAuth)
	social.app = socialAuther

	if len(urlPrefix) == 0 {
		urlPrefix = defaultURLPrefix
	}

	if urlPrefix[len(urlPrefix)-1] != '/' {
		urlPrefix += "/"
	}

	social.URLPrefix = urlPrefix

	social.ConnectSuccessURL = defaultConnectSuccessURL
	social.ConnectFailedURL = defaultConnectFailedURL
	social.LoginURL = defaultLoginURL
	social.ConnectRegisterURL = defaultConnectRegisterURL

	return social
}

// create a instance and create filter
func NewWithFilter(urlPrefix string, socialAuther SocialAuther) *SocialAuth {
	social := NewSocial(urlPrefix, socialAuther)

	beego.InsertFilter(social.URLPrefix+"*/access", beego.BeforeRouter, social.handleAccess)
	beego.InsertFilter(social.URLPrefix+"*", beego.BeforeRouter, social.handleRedirect)

	return social
}
<commit_msg>1）修改默认验证失败后的链接跳转<commit_after>// Copyright 2014 beego 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.
//
// Maintain by https://github.com/slene

package social

import (
	"encoding/base64"
	"fmt"
	"net/url"
	"strconv"

	"github.com/astaxie/beego"
	"github.com/astaxie/beego/context"
	"github.com/astaxie/beego/utils"
)

const (
	defaultURLPrefix          = "/loginsocial/"
	defaultConnectSuccessURL  = "/?flag=connect_success"
	defaultConnectFailedURL   = "/?flag=connect_failed"
	defaultLoginURL           = "/loginsocial"
	defaultConnectRegisterURL = "/register/connect"
)

type SocialAuth struct {
	app                SocialAuther
	URLPrefix          string
	ConnectSuccessURL  string
	ConnectFailedURL   string
	LoginURL           string
	ConnectRegisterURL string
}

// generate session key for social-auth
func (this *SocialAuth) getSessKey(social SocialType, key string) string {
	return "social_" + fmt.Sprintf("%v", social) + "_" + key
}

// create oauth2 state string
func (this *SocialAuth) createState(ctx *context.Context, social SocialType) string {
	values := make(url.Values, 2)

	if uid, ok := this.app.IsUserLogin(ctx); ok {
		// add uid if user current is login
		values.Add("uid", strconv.FormatInt(int64(uid), 10))
	}

	// our secret string
	values.Add("secret", string(utils.RandomCreateBytes(20)))

	// create state string
	state := base64.URLEncoding.EncodeToString([]byte(values.Encode()))

	// save to session
	name := this.getSessKey(social, "state")
	ctx.Input.CruSession.Set(name, state)

	return state
}

// verify oauth2 state string
func (this *SocialAuth) verifyState(ctx *context.Context, social SocialType) (string, bool) {
	code := ctx.Input.Query("code")
	state := ctx.Input.Query("state")

	if len(code) == 0 || len(state) == 0 {
		return "", false
	}

	name := this.getSessKey(social, "state")

	vu, ok := ctx.Input.CruSession.Get(name).(string)
	if !ok || ok && state != vu {
		return "", false
	}

	return code, true
}

// Get provider according request path. ex: /login/: match /login/github
func (this *SocialAuth) getProvider(ctx *context.Context) Provider {
	path := ctx.Input.Param(":splat")

	p, ok := GetProviderByPath(path)
	if ok {
		return p
	}

	return nil
}

// After OAuthAccess check saved token for ready connect
func (this *SocialAuth) ReadyConnect(ctx *context.Context) (SocialType, bool) {
	var social SocialType

	if s, _ := ctx.Input.CruSession.Get("social_connect").(int); s == 0 {
		return 0, false
	} else {
		social = SocialType(s)
	}

	if !social.Available() {
		return 0, false
	}

	return social, true
}

// Redirect to other social platform
func (this *SocialAuth) OAuthRedirect(ctx *context.Context) (redirect string, failedErr error) {
	_, isLogin := this.app.IsUserLogin(ctx)

	defer func() {
		if len(redirect) == 0 && failedErr != nil {
			if isLogin {
				redirect = this.ConnectFailedURL
			} else {
				redirect = this.LoginURL
			}
		}
	}()

	var p Provider
	if p = this.getProvider(ctx); p == nil {
		failedErr = fmt.Errorf("unknown provider")
		return
	}

	social := p.GetType()
	config := p.GetConfig()
	// create redirect url
	redirect = config.AuthCodeURL(this.createState(ctx, social))
	return
}

// Callback from social platform
func (this *SocialAuth) OAuthAccess(ctx *context.Context) (redirect string, userSocial *UserSocial, failedErr error) {
	_, isLogin := this.app.IsUserLogin(ctx)

	defer func() {
		fmt.Println("redirect-:", redirect)
		if len(redirect) == 0 {
			if failedErr != nil {
				if isLogin {
					redirect = this.ConnectFailedURL
				} else {
					redirect = this.LoginURL
				}
			}
		}
	}()

	// check if param has a error key
	if err := ctx.Input.Query("error"); len(err) > 0 {
		failedErr = fmt.Errorf(err)
		return
	}

	// get provider from matched url path
	var p Provider
	if p = this.getProvider(ctx); p == nil {
		failedErr = fmt.Errorf("unknown provider")
		return
	}

	social := p.GetType()

	var code string

	// verify state string
	if c, ok := this.verifyState(ctx, social); !ok {
		failedErr = fmt.Errorf("state not verified")
		return
	} else {
		code = c
	}

	config := p.GetConfig()
	trans := &Transport{config, nil, DefaultTransport}

	// Send code to platform then get token
	if tok, err := trans.Exchange(code); err != nil {
		// get access token
		failedErr = err
	} else if err := tok.GetExtra("error"); err != "" {
		// token has error
		failedErr = fmt.Errorf(err)
	} else if tok.IsEmpty() {
		failedErr = fmt.Errorf("empty access token")
	} else {

		// check
		var uSocial = UserSocial{}
		if ok, err := p.CanConnect(tok, &uSocial); ok {
			// save token to session, for connect
			tk := SocialTokenField{tok}
			ctx.Input.CruSession.Set(this.getSessKey(social, "token"), tk.RawValue())
			ctx.Input.CruSession.Set("social_connect", int(social))

			redirect = this.ConnectRegisterURL
			fmt.Println("redirect-1:", redirect)
		} else if err == nil {
			if !isLogin {
				// login user
				redirect, failedErr = this.app.LoginUser(ctx, uSocial.Uid)
				if len(redirect) == 0 && failedErr == nil {
					redirect = this.ConnectSuccessURL
				}
				fmt.Println("redirect-2:", redirect)
			} else {
				redirect = this.ConnectSuccessURL
				fmt.Println("redirect-3:", redirect)
			}

			// save new access token if it changed
			uSocial.PutToken(tok)

			userSocial = &uSocial

		} else {
			failedErr = err
		}
	}

	return
}

// general use of redirect
func (this *SocialAuth) handleRedirect(ctx *context.Context) {
	redirect, err := this.OAuthRedirect(ctx)
	if err != nil {
		beego.Error("SocialAuth.handleRedirect", err)
	}

	if len(redirect) > 0 {
		ctx.Redirect(302, redirect)
	}
}

// general use of redirect callback
func (this *SocialAuth) handleAccess(ctx *context.Context) {
	redirect, _, err := this.OAuthAccess(ctx)
	if err != nil {
		beego.Error("SocialAuth.handleAccess", err)
	}

	if len(redirect) > 0 {
		ctx.Redirect(302, redirect)
	}
}

// save user social info and login the user
func (this *SocialAuth) ConnectAndLogin(ctx *context.Context, socialType SocialType, uid int64) (string, *UserSocial, error) {
	tokKey := this.getSessKey(socialType, "token")

	defer func() {
		// delete connect tok in session
		if ctx.Input.CruSession.Get("social_connect") != nil {
			ctx.Input.CruSession.Delete("social_connect")
		}
		if ctx.Input.CruSession.Get(tokKey) != nil {
			ctx.Input.CruSession.Delete(tokKey)
		}
	}()

	tk := SocialTokenField{}
	value := ctx.Input.CruSession.Get(tokKey)
	if err := tk.SetRaw(value); err != nil {
		return "", nil, err
	}

	var p Provider
	if p, _ = GetProviderByType(socialType); p == nil {
		return "", nil, fmt.Errorf("unknown provider")
	}

	identify, err := p.GetIndentify(tk.Token)
	if err != nil {
		return "", nil, err
	}
	if len(identify) == 0 {
		return "", nil, fmt.Errorf("empty identify")
	}

	userSocial := UserSocial{
		Uid:      uid,
		Type:     socialType,
		Data:     tk,
		Identify: identify,
	}

	if err := userSocial.Save(); err != nil {
		return "", nil, err
	}

	// login user
	loginRedirect, err := this.app.LoginUser(ctx, uid)
	return loginRedirect, &userSocial, nil
}

// create a global SocialAuth instance
func NewSocial(urlPrefix string, socialAuther SocialAuther) *SocialAuth {
	social := new(SocialAuth)
	social.app = socialAuther

	if len(urlPrefix) == 0 {
		urlPrefix = defaultURLPrefix
	}

	if urlPrefix[len(urlPrefix)-1] != '/' {
		urlPrefix += "/"
	}

	social.URLPrefix = urlPrefix

	social.ConnectSuccessURL = defaultConnectSuccessURL
	social.ConnectFailedURL = defaultConnectFailedURL
	social.LoginURL = defaultLoginURL
	social.ConnectRegisterURL = defaultConnectRegisterURL

	return social
}

// create a instance and create filter
func NewWithFilter(urlPrefix string, socialAuther SocialAuther) *SocialAuth {
	social := NewSocial(urlPrefix, socialAuther)

	beego.InsertFilter(social.URLPrefix+"*/access", beego.BeforeRouter, social.handleAccess)
	beego.InsertFilter(social.URLPrefix+"*", beego.BeforeRouter, social.handleRedirect)

	return social
}
<|endoftext|>


<commit_before>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"errors"
	"net"
	"time"
)

type Socket struct {
	pc            net.PacketConn
	ctx           *C.utp_context
	backlog       chan *Conn
	closed        bool
	conns         map[*C.utp_socket]*Conn
	nonUtpReads   chan packet
	writeDeadline time.Time
	readDeadline  time.Time
}

var (
	_ net.PacketConn = (*Socket)(nil)
)

type packet struct {
	b    []byte
	from net.Addr
}

func NewSocket(network, addr string) (*Socket, error) {
	pc, err := net.ListenPacket(network, addr)
	if err != nil {
		return nil, err
	}
	mu.Lock()
	defer mu.Unlock()
	ctx := C.utp_init(2)
	if ctx == nil {
		panic(ctx)
	}
	ctx.setCallbacks()
	// ctx.setOption(C.UTP_LOG_NORMAL, 1)
	// ctx.setOption(C.UTP_LOG_MTU, 1)
	// ctx.setOption(C.UTP_LOG_DEBUG, 1)
	s := &Socket{
		pc:      pc,
		ctx:     ctx,
		backlog: make(chan *Conn, 5),
		conns:   make(map[*C.utp_socket]*Conn),
	}
	libContextToSocket[ctx] = s
	go s.packetReader()
	return s, nil
}

func (s *Socket) onLibSocketDestroyed(ls *C.utp_socket) {
	delete(s.conns, ls)
}

func (s *Socket) newConn(us *C.utp_socket) *Conn {
	c := &Conn{
		s: us,
	}
	c.cond.L = &mu
	s.conns[us] = c
	return c
}

func (s *Socket) packetReader() {
	for {
		b := make([]byte, 0x1000)
		n, addr, err := s.pc.ReadFrom(b)
		if err != nil {
			mu.Lock()
			closed := s.closed
			mu.Unlock()
			if closed {
				return
			}
			panic(err)
		}
		sa, sal := netAddrToLibSockaddr(addr)
		func() {
			mu.Lock()
			defer mu.Unlock()
			if s.closed {
				return
			}
			if C.utp_process_udp(s.ctx, (*C.byte)(&b[0]), C.size_t(n), sa, sal) == 0 {
				s.onReadNonUtp(b[:n], addr)
			}
			C.utp_issue_deferred_acks(s.ctx)
			C.utp_check_timeouts(s.ctx)
		}()
	}
}

func (s *Socket) timeoutChecker() {
	for {
		mu.Lock()
		if s.closed {
			mu.Unlock()
			return
		}
		C.utp_check_timeouts(s.ctx)
		mu.Unlock()
		time.Sleep(500 * time.Millisecond)
	}
}

func (me *Socket) Close() error {
	mu.Lock()
	defer mu.Unlock()
	if me.closed {
		return nil
	}
	C.utp_destroy(me.ctx)
	me.ctx = nil
	me.pc.Close()
	close(me.backlog)
	me.closed = true
	return nil
}

func (me *Socket) Addr() net.Addr {
	return me.pc.LocalAddr()
}

func (me *Socket) LocalAddr() net.Addr {
	return me.pc.LocalAddr()
}

func (s *Socket) Accept() (net.Conn, error) {
	nc, ok := <-s.backlog
	if !ok {
		return nil, errors.New("closed")
	}
	return nc, nil
}

func (s *Socket) Dial(addr string) (net.Conn, error) {
	return s.DialTimeout(addr, 0)
}

func (s *Socket) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	ua, err := net.ResolveUDPAddr(s.Addr().Network(), addr)
	if err != nil {
		panic(err)
	}
	sa, sl := netAddrToLibSockaddr(ua)
	mu.Lock()
	c := s.newConn(C.utp_create_socket(s.ctx))
	C.utp_connect(c.s, sa, sl)
	defer mu.Unlock()
	err = c.waitForConnect()
	if err != nil {
		c.close()
		return nil, err
	}
	return c, nil
}

func (s *Socket) pushBacklog(c *Conn) {
	select {
	case s.backlog <- c:
	default:
		c.Close()
	}
}

func (s *Socket) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	p := <-s.nonUtpReads
	n = copy(b, p.b)
	addr = p.from
	return
}

func (s *Socket) onReadNonUtp(b []byte, from net.Addr) {
	select {
	case s.nonUtpReads <- packet{b, from}:
	default:
	}
}

func (s *Socket) SetReadDeadline(t time.Time) error {
	return nil
}

func (s *Socket) SetWriteDeadline(t time.Time) error {
	return nil
}

func (s *Socket) SetDeadline(t time.Time) error {
	return nil
}

func (s *Socket) WriteTo(b []byte, addr net.Addr) (int, error) {
	return s.pc.WriteTo(b, addr)
}
<commit_msg>Fix non uTP read buffer<commit_after>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"errors"
	"log"
	"net"
	"time"
)

type Socket struct {
	pc            net.PacketConn
	ctx           *C.utp_context
	backlog       chan *Conn
	closed        bool
	conns         map[*C.utp_socket]*Conn
	nonUtpReads   chan packet
	writeDeadline time.Time
	readDeadline  time.Time
}

var (
	_ net.PacketConn = (*Socket)(nil)
)

type packet struct {
	b    []byte
	from net.Addr
}

func NewSocket(network, addr string) (*Socket, error) {
	pc, err := net.ListenPacket(network, addr)
	if err != nil {
		return nil, err
	}
	mu.Lock()
	defer mu.Unlock()
	ctx := C.utp_init(2)
	if ctx == nil {
		panic(ctx)
	}
	ctx.setCallbacks()
	// ctx.setOption(C.UTP_LOG_NORMAL, 1)
	// ctx.setOption(C.UTP_LOG_MTU, 1)
	// ctx.setOption(C.UTP_LOG_DEBUG, 1)
	s := &Socket{
		pc:          pc,
		ctx:         ctx,
		backlog:     make(chan *Conn, 5),
		conns:       make(map[*C.utp_socket]*Conn),
		nonUtpReads: make(chan packet, 10),
	}
	libContextToSocket[ctx] = s
	go s.packetReader()
	return s, nil
}

func (s *Socket) onLibSocketDestroyed(ls *C.utp_socket) {
	delete(s.conns, ls)
}

func (s *Socket) newConn(us *C.utp_socket) *Conn {
	c := &Conn{
		s: us,
	}
	c.cond.L = &mu
	s.conns[us] = c
	return c
}

func (s *Socket) packetReader() {
	for {
		b := make([]byte, 0x1000)
		n, addr, err := s.pc.ReadFrom(b)
		if err != nil {
			mu.Lock()
			closed := s.closed
			mu.Unlock()
			if closed {
				return
			}
			panic(err)
		}
		sa, sal := netAddrToLibSockaddr(addr)
		func() {
			mu.Lock()
			defer mu.Unlock()
			if s.closed {
				return
			}
			if C.utp_process_udp(s.ctx, (*C.byte)(&b[0]), C.size_t(n), sa, sal) == 0 {
				s.onReadNonUtp(b[:n], addr)
			}
			C.utp_issue_deferred_acks(s.ctx)
			C.utp_check_timeouts(s.ctx)
		}()
	}
}

func (s *Socket) timeoutChecker() {
	for {
		mu.Lock()
		if s.closed {
			mu.Unlock()
			return
		}
		C.utp_check_timeouts(s.ctx)
		mu.Unlock()
		time.Sleep(500 * time.Millisecond)
	}
}

func (me *Socket) Close() error {
	mu.Lock()
	defer mu.Unlock()
	if me.closed {
		return nil
	}
	C.utp_destroy(me.ctx)
	me.ctx = nil
	me.pc.Close()
	close(me.backlog)
	me.closed = true
	return nil
}

func (me *Socket) Addr() net.Addr {
	return me.pc.LocalAddr()
}

func (me *Socket) LocalAddr() net.Addr {
	return me.pc.LocalAddr()
}

func (s *Socket) Accept() (net.Conn, error) {
	nc, ok := <-s.backlog
	if !ok {
		return nil, errors.New("closed")
	}
	return nc, nil
}

func (s *Socket) Dial(addr string) (net.Conn, error) {
	return s.DialTimeout(addr, 0)
}

func (s *Socket) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	ua, err := net.ResolveUDPAddr(s.Addr().Network(), addr)
	if err != nil {
		panic(err)
	}
	sa, sl := netAddrToLibSockaddr(ua)
	mu.Lock()
	c := s.newConn(C.utp_create_socket(s.ctx))
	C.utp_connect(c.s, sa, sl)
	defer mu.Unlock()
	err = c.waitForConnect()
	if err != nil {
		c.close()
		return nil, err
	}
	return c, nil
}

func (s *Socket) pushBacklog(c *Conn) {
	select {
	case s.backlog <- c:
	default:
		c.Close()
	}
}

func (s *Socket) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	p := <-s.nonUtpReads
	n = copy(b, p.b)
	addr = p.from
	return
}

func (s *Socket) onReadNonUtp(b []byte, from net.Addr) {
	select {
	case s.nonUtpReads <- packet{b, from}:
	default:
		log.Printf("dropped non utp packet: no room in buffer")
	}
}

func (s *Socket) SetReadDeadline(t time.Time) error {
	return nil
}

func (s *Socket) SetWriteDeadline(t time.Time) error {
	return nil
}

func (s *Socket) SetDeadline(t time.Time) error {
	return nil
}

func (s *Socket) WriteTo(b []byte, addr net.Addr) (int, error) {
	return s.pc.WriteTo(b, addr)
}
<|endoftext|>


<commit_before>package socks5

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"

	"github.com/hmgle/tcfs-go"
)

const (
	socks5Version = uint8(5)
)

const (
	ipv4Address = uint8(1)
	fqdnAddress = uint8(3)
	ipv6Address = uint8(4)
)

var (
	errBadAddrType = fmt.Errorf("Unrecognized address type")
)

// AddrSpec is used to return the target AddrSpec
type AddrSpec struct {
	Addr string
	Port int
}

type Local struct {
	faddr, baddr *net.TCPAddr
	cipher       *tcfs.Cipher
}

type Server struct {
	server *net.TCPAddr
	cipher *tcfs.Cipher
}

func NewLocal(faddr, baddr string, cryptoMethod string, key []byte) *Local {
	a1, err := net.ResolveTCPAddr("tcp", faddr)
	if err != nil {
		log.Fatalln("resolve frontend error:", err)
	}
	a2, err := net.ResolveTCPAddr("tcp", baddr)
	if err != nil {
		log.Fatalln("resolve backend error:", err)
	}
	var cipher *tcfs.Cipher
	if len(key) > 0 {
		cipher = tcfs.NewCipher(cryptoMethod, key)
	}
	return &Local{a1, a2, cipher}
}

func NewServer(port string, cryptoMethod string, key []byte) *Server {
	addr, err := net.ResolveTCPAddr("tcp", port)
	if err != nil {
		log.Fatalln("resolve frontend error:", err)
	}
	var cipher *tcfs.Cipher
	if len(key) > 0 {
		cipher = tcfs.NewCipher(cryptoMethod, key)
	}
	return &Server{addr, cipher}
}

func (s *Local) Start() {
	ln, err := net.ListenTCP("tcp", s.faddr)
	if err != nil {
		log.Fatal(err)
	}
	defer ln.Close()

	for {
		conn, err := ln.AcceptTCP()
		if err != nil {
			log.Println("accept:", err)
			continue
		}
		go s.handleConn(conn)
	}
}

func (s *Local) handleConn(conn net.Conn) error {
	bufConn := bufio.NewReader(conn)

	version := []byte{0}
	if _, err := bufConn.Read(version); err != nil {
		log.Println("read:", err)
		return err
	}
	if version[0] != socks5Version {
		err := fmt.Errorf("Unsupported SOCKS version: %v", version)
		log.Println("[ERR] socks: %v", err)
		return err
	}
	nmethods := []byte{0}
	if _, err := bufConn.Read(nmethods); err != nil {
		log.Println("read:", err)
		return err
	}
	numMethods := int(nmethods[0])
	_methods := make([]byte, numMethods)
	_, err := io.ReadAtLeast(bufConn, _methods, numMethods)
	if err != nil {
		log.Println("ReadAtLeast:", err)
		return err
	}
	conn.Write([]byte{0x05, 0x00})

	backConn, err := net.DialTCP("tcp", nil, s.baddr)
	if err != nil {
		log.Println(err)
		return err
	}
	lConn := NewConn(conn, nil)
	sConn := NewConn(backConn, s.cipher)
	go pipe(sConn, lConn)
	go pipe(lConn, sConn)
	return nil
}

func (s *Server) Start() {
	ln, err := net.ListenTCP("tcp", s.server)
	if err != nil {
		log.Fatal(err)
	}
	defer ln.Close()

	for {
		conn, err := ln.AcceptTCP()
		if err != nil {
			log.Println("accept:", err)
			continue
		}
		go s.handleConn(conn)
	}
}

func (s *Server) handleConn(conn net.Conn) error {
	cConn := NewConn(conn, s.cipher)
	bufConn := bufio.NewReader(cConn)

	header := []byte{0, 0, 0}
	if _, err := io.ReadAtLeast(bufConn, header, 3); err != nil {
		log.Println("io.ReadAtLeast:", err)
		return fmt.Errorf("Failed to get command version: %v", err)
	}

	// Ensure we are compatible
	if header[0] != socks5Version {
		log.Println("bad version")
		return fmt.Errorf("Unsupported command version: %v", header[0])
	}

	// Read in the destination address
	dest, err := readAddrSpec(bufConn)
	if err != nil {
		log.Println("bad addr")
		return fmt.Errorf("Failed to read dest addr: %v", err)
	}
	cConn.Write([]byte{0x05, 0, 0, 0x01, 0, 0, 0, 0, 0x19, 0x19})
	remote, err := net.Dial("tcp", fmt.Sprintf("%s:%d", dest.Addr, dest.Port))
	if err != nil {
		log.Println("net.Dial failed:", dest.Addr, dest.Port)
		return err
	}

	rConn := NewConn(remote, nil)
	go pipe(rConn, cConn)
	go pipe(cConn, rConn)
	return nil
}

// readAddrSpec is used to read AddrSpec.
// Expects an address type byte, follwed by the address and port
func readAddrSpec(r io.Reader) (*AddrSpec, error) {
	d := &AddrSpec{}

	// Get the address type
	addrType := []byte{0}
	if _, err := r.Read(addrType); err != nil {
		return nil, err
	}

	// Handle on a per type basis
	switch addrType[0] {
	case ipv4Address:
		addr := make([]byte, 4)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.Addr = net.IP(addr).String()
	case ipv6Address:
		addr := make([]byte, 16)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.Addr = net.IP(addr).String()
	case fqdnAddress:
		if _, err := r.Read(addrType); err != nil {
			return nil, err
		}
		addrLen := int(addrType[0])
		fqdn := make([]byte, addrLen)
		if _, err := io.ReadAtLeast(r, fqdn, addrLen); err != nil {
			return nil, err
		}
		d.Addr = string(fqdn)
	default:
		return nil, errBadAddrType
	}

	// Read the port
	port := []byte{0, 0}
	if _, err := io.ReadAtLeast(r, port, 2); err != nil {
		return nil, err
	}
	d.Port = (int(port[0]) << 8) | int(port[1])

	return d, nil
}

func pipe(dst, src *Conn) {
	for {
		_, err := io.Copy(dst, src)
		if err != nil {
			return
		}
	}
}
<commit_msg>update pipe<commit_after>package socks5

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"

	"github.com/hmgle/tcfs-go"
)

const (
	socks5Version = uint8(5)
)

const (
	ipv4Address = uint8(1)
	fqdnAddress = uint8(3)
	ipv6Address = uint8(4)
)

var (
	errBadAddrType = fmt.Errorf("Unrecognized address type")
)

// AddrSpec is used to return the target AddrSpec
type AddrSpec struct {
	Addr string
	Port int
}

type Local struct {
	faddr, baddr *net.TCPAddr
	cipher       *tcfs.Cipher
}

type Server struct {
	server *net.TCPAddr
	cipher *tcfs.Cipher
}

func NewLocal(faddr, baddr string, cryptoMethod string, key []byte) *Local {
	a1, err := net.ResolveTCPAddr("tcp", faddr)
	if err != nil {
		log.Fatalln("resolve frontend error:", err)
	}
	a2, err := net.ResolveTCPAddr("tcp", baddr)
	if err != nil {
		log.Fatalln("resolve backend error:", err)
	}
	var cipher *tcfs.Cipher
	if len(key) > 0 {
		cipher = tcfs.NewCipher(cryptoMethod, key)
	}
	return &Local{a1, a2, cipher}
}

func NewServer(port string, cryptoMethod string, key []byte) *Server {
	addr, err := net.ResolveTCPAddr("tcp", port)
	if err != nil {
		log.Fatalln("resolve frontend error:", err)
	}
	var cipher *tcfs.Cipher
	if len(key) > 0 {
		cipher = tcfs.NewCipher(cryptoMethod, key)
	}
	return &Server{addr, cipher}
}

func (s *Local) Start() {
	ln, err := net.ListenTCP("tcp", s.faddr)
	if err != nil {
		log.Fatal(err)
	}
	defer ln.Close()

	for {
		conn, err := ln.AcceptTCP()
		if err != nil {
			log.Println("accept:", err)
			continue
		}
		go s.handleConn(conn)
	}
}

func (s *Local) handleConn(conn net.Conn) error {
	bufConn := bufio.NewReader(conn)

	version := []byte{0}
	if _, err := bufConn.Read(version); err != nil {
		log.Println("read:", err)
		return err
	}
	if version[0] != socks5Version {
		err := fmt.Errorf("Unsupported SOCKS version: %v", version)
		log.Println("[ERR] socks: %v", err)
		return err
	}
	nmethods := []byte{0}
	if _, err := bufConn.Read(nmethods); err != nil {
		log.Println("read:", err)
		return err
	}
	numMethods := int(nmethods[0])
	_methods := make([]byte, numMethods)
	_, err := io.ReadAtLeast(bufConn, _methods, numMethods)
	if err != nil {
		log.Println("ReadAtLeast:", err)
		return err
	}
	conn.Write([]byte{0x05, 0x00})

	backConn, err := net.DialTCP("tcp", nil, s.baddr)
	if err != nil {
		log.Println(err)
		return err
	}
	lConn := NewConn(conn, nil)
	sConn := NewConn(backConn, s.cipher)
	go pipe(sConn, lConn)
	go pipe(lConn, sConn)
	return nil
}

func (s *Server) Start() {
	ln, err := net.ListenTCP("tcp", s.server)
	if err != nil {
		log.Fatal(err)
	}
	defer ln.Close()

	for {
		conn, err := ln.AcceptTCP()
		if err != nil {
			log.Println("accept:", err)
			continue
		}
		go s.handleConn(conn)
	}
}

func (s *Server) handleConn(conn net.Conn) error {
	cConn := NewConn(conn, s.cipher)
	bufConn := bufio.NewReader(cConn)

	header := []byte{0, 0, 0}
	if _, err := io.ReadAtLeast(bufConn, header, 3); err != nil {
		log.Println("io.ReadAtLeast:", err)
		return fmt.Errorf("Failed to get command version: %v", err)
	}

	// Ensure we are compatible
	if header[0] != socks5Version {
		log.Println("bad version")
		return fmt.Errorf("Unsupported command version: %v", header[0])
	}

	// Read in the destination address
	dest, err := readAddrSpec(bufConn)
	if err != nil {
		log.Println("bad addr")
		return fmt.Errorf("Failed to read dest addr: %v", err)
	}
	cConn.Write([]byte{0x05, 0, 0, 0x01, 0, 0, 0, 0, 0x19, 0x19})
	remote, err := net.Dial("tcp", fmt.Sprintf("%s:%d", dest.Addr, dest.Port))
	if err != nil {
		log.Println("net.Dial failed:", dest.Addr, dest.Port)
		return err
	}

	rConn := NewConn(remote, nil)
	go pipe(rConn, cConn)
	go pipe(cConn, rConn)
	return nil
}

// readAddrSpec is used to read AddrSpec.
// Expects an address type byte, follwed by the address and port
func readAddrSpec(r io.Reader) (*AddrSpec, error) {
	d := &AddrSpec{}

	// Get the address type
	addrType := []byte{0}
	if _, err := r.Read(addrType); err != nil {
		return nil, err
	}

	// Handle on a per type basis
	switch addrType[0] {
	case ipv4Address:
		addr := make([]byte, 4)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.Addr = net.IP(addr).String()
	case ipv6Address:
		addr := make([]byte, 16)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.Addr = net.IP(addr).String()
	case fqdnAddress:
		if _, err := r.Read(addrType); err != nil {
			return nil, err
		}
		addrLen := int(addrType[0])
		fqdn := make([]byte, addrLen)
		if _, err := io.ReadAtLeast(r, fqdn, addrLen); err != nil {
			return nil, err
		}
		d.Addr = string(fqdn)
	default:
		return nil, errBadAddrType
	}

	// Read the port
	port := []byte{0, 0}
	if _, err := io.ReadAtLeast(r, port, 2); err != nil {
		return nil, err
	}
	d.Port = (int(port[0]) << 8) | int(port[1])

	return d, nil
}

func pipe(dst, src *Conn) {
	io.Copy(dst, src)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"io"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

type Source struct {
	sinkCount       int
	frames          []DataFrame
	frameLocks      []sync.RWMutex
	frameBytes      uint64
	gone            bool
	header          []byte
	HeaderBytes     uint64
	input           io.ReadCloser
	nextFrame       uint64 // How many frames have ever been here
	path            string
	closeIdle       bool
	reopen          bool
	bandwidth       uint64
	sync.RWMutex    // Must be held while changing nextFrame or gone
	*sync.Cond      // Control access to frames other than nextFrame
	statBytesIn     uint64
	statBytesOut    uint64
	statLast        time.Time
	statLogInterval time.Duration
}

var sourceMap = make(map[string]*Source)
var sourceMapMutex = sync.RWMutex{}

func NewSource(path string, c *Config) (s *Source) {
	s = &Source{}
	s.Cond = sync.NewCond(s.RLocker())
	s.frameLocks = make([]sync.RWMutex, c.SourceBuffer)
	s.frames = make([]DataFrame, c.SourceBuffer)
	for i := range s.frames {
		s.frames[i] = make(DataFrame, c.FrameBytes)
	}
	s.bandwidth = c.SourceBandwidth
	s.closeIdle = c.CloseIdle
	s.frameBytes = c.FrameBytes
	s.HeaderBytes = c.HeaderBytes
	s.path = path
	s.reopen = c.Reopen
	s.statLogInterval = c.StatLogInterval
	return
}

func (s *Source) openInput() (err error) {
	// Notify anyone waiting for the header to arrive
	defer s.Cond.Broadcast()
	if s.input, err = os.Open(s.path); err != nil {
		log.Printf("source %s open: %s", s.path, err)
		return
	}
	header := make([]byte, s.HeaderBytes)
	for pos := uint64(0); pos < s.HeaderBytes; {
		var got int
		if got, err = s.input.Read(header[pos:]); got == 0 {
			log.Printf("source %s read-header: %s", s.path, err)
			s.input.Close()
			return
		}
		pos += uint64(got)
	}
	if len(s.header) > 0 && bytes.Compare(header, s.header) != 0 {
		log.Printf("header mismatch: old %v, new %v", s.header, header)
		s.input.Close()
		return
	}
	s.Cond.L.Lock()
	s.header = header
	s.Cond.L.Unlock()
	s.statBytesIn += s.HeaderBytes
	return
}

func (s *Source) readNextFrame() (err error) {
	bufPos := s.nextFrame % uint64(cap(s.frames))
	s.frameLocks[bufPos].Lock()
	defer s.frameLocks[bufPos].Unlock()
	for framePos := uint64(0); framePos < s.frameBytes; {
		var got int
		if got, err = s.input.Read(s.frames[bufPos][framePos:]); got <= 0 {
			return
		}
		framePos += uint64(got)
	}
	return
}

// Read data from the given source into the buffer. If the source
// reaches EOF and cannot be reopened, remove the source from
// sourceMap and return.
func (s *Source) run() {
	var err error
	s.statLast = time.Now()
	defer s.LogStats(true)
	defer s.Close()
	if err := s.openInput(); err != nil {
		return
	}
	var ticker *time.Ticker
	if s.bandwidth > 0 {
		ticker = time.NewTicker(time.Duration(uint64(1000000000) * s.frameBytes / s.bandwidth))
	}
	for !s.gone {
		if err = s.readNextFrame(); err != nil {
			log.Printf("source %s read: %s", s.path, err)
			s.input.Close()
			s.input = nil
			if !s.reopen {
				break
			} else if err = s.openInput(); err != nil {
				continue
			} else {
				break
			}
		}
		s.nextFrame += 1
		s.Cond.Broadcast()
		s.statBytesIn += s.frameBytes
		s.LogStats(false)
		if ticker != nil {
			<-ticker.C
		}
	}
	if s.input != nil {
		s.input.Close()
	}
}

// If !really, only if statLogInterval says so.
func (s *Source) LogStats(really bool) {
	if really || (s.statLogInterval > 0 && time.Since(s.statLast) >= s.statLogInterval) {
		s.statLast.Add(s.statLogInterval)
		log.Printf("source %s stats: %d in %d out", s.path, s.statBytesIn, s.statBytesOut)
	}
}

func (s *Source) GetHeader(buf []byte) (err error) {
	s.Cond.L.Lock()
	defer s.Cond.L.Unlock()
	for uint64(len(s.header)) < s.HeaderBytes && !s.gone {
		s.Cond.Wait()
	}
	if uint64(len(s.header)) < s.HeaderBytes {
		err = io.EOF
		return
	}
	if cap(buf) < len(s.header) {
		err = errors.New("Caller's header buffer is too small.")
		return
	}
	atomic.AddUint64(&s.statBytesOut, s.HeaderBytes)
	copy(buf, s.header)
	return
}

// Copy the next data frame into the given buffer and update the
// nextFrame pointer.
//
// Return the number of frames skipped due to buffer underrun. If the
// data source is exhausted, return with err != nil (with frame
// untouched and other return values undefined).
func (s *Source) Next(nextFrame *uint64, frame DataFrame) (nSkipped uint64, err error) {
	defer func() {
		if err == nil {
			atomic.AddUint64(&s.statBytesOut, s.frameBytes)
			*nextFrame += 1
		}
	}()
	if s.nextFrame > uint64(1) && *nextFrame == uint64(0) {
		// New clients start out reading fresh frames.
		*nextFrame = s.nextFrame - uint64(1)
	}
	s.Cond.L.Lock()
	for *nextFrame >= s.nextFrame && !s.gone {
		s.Cond.Wait()
	}
	s.Cond.L.Unlock()
	if *nextFrame >= s.nextFrame {
		err = io.EOF
		return
	}
	if s.nextFrame >= *nextFrame + uint64(cap(s.frames)) {
		// s.nextFrame has lapped *nextFrame. Catch up.
		delta := s.nextFrame - *nextFrame - uint64(1)
		nSkipped += delta
		*nextFrame += delta
	}
	bufPos := *nextFrame % uint64(cap(s.frames))
	s.frameLocks[bufPos].RLock()
	defer s.frameLocks[bufPos].RUnlock()
	copy(frame, s.frames[bufPos])
	if cap(frame) < len(s.frames[bufPos]) {
		err = errors.New("Caller's frame buffer is too small.")
	}
	return
}

func (s *Source) Done() {
	s.sinkCount -= 1
	if s.closeIdle {
		s.CloseIfIdle()
	}
}

// Make sure everyone waiting in Next() gives up. Prevents deadlock.
func (s *Source) disconnectAll() {
	s.gone = true
	s.Broadcast()
}

func (s *Source) Close() {
	s.closeIdle = true
	s.disconnectAll()
}

func (s *Source) CloseIfIdle() {
	didClose := false
	sourceMapMutex.Lock()
	if s.sinkCount == 0 {
		delete(sourceMap, s.path)
		didClose = true
	}
	sourceMapMutex.Unlock()
	if didClose {
		s.disconnectAll()
	}
}

func CloseAllSources() {
	for _, src := range sourceMap {
		src.Close()
	}
}

// Return a Source for the given path (URI) and config (argv). At any
// given time, there is at most one Source for a given path.
//
// The caller must ensure Done() is eventually called, exactly once,
// on the returned *Source.
func GetSource(path string, c *Config) (src *Source) {
	sourceMapMutex.Lock()
	defer sourceMapMutex.Unlock()
	var ok bool
	if src, ok = sourceMap[path]; !ok {
		src = NewSource(path, c)
		sourceMap[path] = src
		go src.run()
	}
	src.sinkCount += 1
	return src
}
<commit_msg>Skip some unnecessary locking. Always try to catch sink up to fresh frames if source has lapped it.<commit_after>package main

import (
	"bytes"
	"errors"
	"io"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

type Source struct {
	sinkCount       int
	frames          []DataFrame
	frameLocks      []sync.RWMutex
	frameBytes      uint64
	gone            bool
	header          []byte
	HeaderBytes     uint64
	input           io.ReadCloser
	nextFrame       uint64 // How many frames have ever been here
	path            string
	closeIdle       bool
	reopen          bool
	bandwidth       uint64
	sync.RWMutex    // Must be held while changing nextFrame or gone
	*sync.Cond      // Control access to frames other than nextFrame
	statBytesIn     uint64
	statBytesOut    uint64
	statLast        time.Time
	statLogInterval time.Duration
}

var sourceMap = make(map[string]*Source)
var sourceMapMutex = sync.RWMutex{}

func NewSource(path string, c *Config) (s *Source) {
	s = &Source{}
	s.Cond = sync.NewCond(s.RLocker())
	s.frameLocks = make([]sync.RWMutex, c.SourceBuffer)
	s.frames = make([]DataFrame, c.SourceBuffer)
	for i := range s.frames {
		s.frames[i] = make(DataFrame, c.FrameBytes)
	}
	s.bandwidth = c.SourceBandwidth
	s.closeIdle = c.CloseIdle
	s.frameBytes = c.FrameBytes
	s.HeaderBytes = c.HeaderBytes
	s.path = path
	s.reopen = c.Reopen
	s.statLogInterval = c.StatLogInterval
	return
}

func (s *Source) openInput() (err error) {
	// Notify anyone waiting for the header to arrive
	defer s.Cond.Broadcast()
	if s.input, err = os.Open(s.path); err != nil {
		log.Printf("source %s open: %s", s.path, err)
		return
	}
	header := make([]byte, s.HeaderBytes)
	for pos := uint64(0); pos < s.HeaderBytes; {
		var got int
		if got, err = s.input.Read(header[pos:]); got == 0 {
			log.Printf("source %s read-header: %s", s.path, err)
			s.input.Close()
			return
		}
		pos += uint64(got)
	}
	if len(s.header) > 0 && bytes.Compare(header, s.header) != 0 {
		log.Printf("header mismatch: old %v, new %v", s.header, header)
		s.input.Close()
		return
	}
	s.Cond.L.Lock()
	s.header = header
	s.Cond.L.Unlock()
	s.statBytesIn += s.HeaderBytes
	return
}

func (s *Source) readNextFrame() (err error) {
	bufPos := s.nextFrame % uint64(cap(s.frames))
	s.frameLocks[bufPos].Lock()
	defer s.frameLocks[bufPos].Unlock()
	for framePos := uint64(0); framePos < s.frameBytes; {
		var got int
		if got, err = s.input.Read(s.frames[bufPos][framePos:]); got <= 0 {
			return
		}
		framePos += uint64(got)
	}
	return
}

// Read data from the given source into the buffer. If the source
// reaches EOF and cannot be reopened, remove the source from
// sourceMap and return.
func (s *Source) run() {
	var err error
	s.statLast = time.Now()
	defer s.LogStats(true)
	defer s.Close()
	if err := s.openInput(); err != nil {
		return
	}
	var ticker *time.Ticker
	if s.bandwidth > 0 {
		ticker = time.NewTicker(time.Duration(uint64(1000000000) * s.frameBytes / s.bandwidth))
	}
	for !s.gone {
		if err = s.readNextFrame(); err != nil {
			log.Printf("source %s read: %s", s.path, err)
			s.input.Close()
			s.input = nil
			if !s.reopen {
				break
			} else if err = s.openInput(); err != nil {
				continue
			} else {
				break
			}
		}
		s.nextFrame += 1
		s.Cond.Broadcast()
		s.statBytesIn += s.frameBytes
		s.LogStats(false)
		if ticker != nil {
			<-ticker.C
		}
	}
	if s.input != nil {
		s.input.Close()
	}
}

// If !really, only if statLogInterval says so.
func (s *Source) LogStats(really bool) {
	if really || (s.statLogInterval > 0 && time.Since(s.statLast) >= s.statLogInterval) {
		s.statLast.Add(s.statLogInterval)
		log.Printf("source %s stats: %d in %d out", s.path, s.statBytesIn, s.statBytesOut)
	}
}

func (s *Source) GetHeader(buf []byte) (err error) {
	s.Cond.L.Lock()
	defer s.Cond.L.Unlock()
	for uint64(len(s.header)) < s.HeaderBytes && !s.gone {
		s.Cond.Wait()
	}
	if uint64(len(s.header)) < s.HeaderBytes {
		err = io.EOF
		return
	}
	if cap(buf) < len(s.header) {
		err = errors.New("Caller's header buffer is too small.")
		return
	}
	atomic.AddUint64(&s.statBytesOut, s.HeaderBytes)
	copy(buf, s.header)
	return
}

// Copy the next data frame into the given buffer and update the
// nextFrame pointer.
//
// Return the number of frames skipped due to buffer underrun. If the
// data source is exhausted, return with err != nil (with frame
// untouched and other return values undefined).
func (s *Source) Next(nextFrame *uint64, frame DataFrame) (nSkipped uint64, err error) {
	defer func() {
		if err == nil {
			atomic.AddUint64(&s.statBytesOut, s.frameBytes)
			*nextFrame += 1
		}
	}()
	if s.nextFrame > uint64(1) && *nextFrame == uint64(0) {
		// New clients start out reading fresh frames.
		*nextFrame = s.nextFrame - uint64(1)
	} else if s.nextFrame >= *nextFrame + uint64(cap(s.frames)) {
		// s.nextFrame has lapped *nextFrame. Catch up.
		delta := s.nextFrame - *nextFrame - uint64(1)
		nSkipped += delta
		*nextFrame += delta
	} else if *nextFrame >= s.nextFrame {
		s.Cond.L.Lock()
		for *nextFrame >= s.nextFrame && !s.gone {
			s.Cond.Wait()
		}
		s.Cond.L.Unlock()
		if *nextFrame >= s.nextFrame {
			err = io.EOF
			return
		}
	}
	bufPos := *nextFrame % uint64(cap(s.frames))
	s.frameLocks[bufPos].RLock()
	defer s.frameLocks[bufPos].RUnlock()
	copy(frame, s.frames[bufPos])
	if cap(frame) < len(s.frames[bufPos]) {
		err = errors.New("Caller's frame buffer is too small.")
	}
	return
}

func (s *Source) Done() {
	s.sinkCount -= 1
	if s.closeIdle {
		s.CloseIfIdle()
	}
}

// Make sure everyone waiting in Next() gives up. Prevents deadlock.
func (s *Source) disconnectAll() {
	s.gone = true
	s.Broadcast()
}

func (s *Source) Close() {
	s.closeIdle = true
	s.disconnectAll()
}

func (s *Source) CloseIfIdle() {
	didClose := false
	sourceMapMutex.Lock()
	if s.sinkCount == 0 {
		delete(sourceMap, s.path)
		didClose = true
	}
	sourceMapMutex.Unlock()
	if didClose {
		s.disconnectAll()
	}
}

func CloseAllSources() {
	for _, src := range sourceMap {
		src.Close()
	}
}

// Return a Source for the given path (URI) and config (argv). At any
// given time, there is at most one Source for a given path.
//
// The caller must ensure Done() is eventually called, exactly once,
// on the returned *Source.
func GetSource(path string, c *Config) (src *Source) {
	sourceMapMutex.Lock()
	defer sourceMapMutex.Unlock()
	var ok bool
	if src, ok = sourceMap[path]; !ok {
		src = NewSource(path, c)
		sourceMap[path] = src
		go src.run()
	}
	src.sinkCount += 1
	return src
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"math/big"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
)

type LnEndpoint string
type LnAddr btcutil.Address
type NodePubKey *btcec.PublicKey

type SharedSecret wire.ShaHash

var zeroNode wire.ShaHash

type MixHeader []byte

// SphinxHeader...
type SphinxHeader struct {
}

// GenerateSphinxHeader...
// TODO(roasbeef): or pass in identifiers as payment path? have map from id -> pubkey
func GenerateSphinxHeader(dest LnAddr, identifier wire.ShaHash, paymentPath []NodePubKey) (*MixHeader, []SharedSecret, error) {
	// Each hop performs ECDH with our ephemeral key pair to arrive at a
	// shared secret. Additionally, each hop randomizes the group element
	// for the next hop by multiplying it by the blinding factor. This way
	// we only need to transmit a single group element, and hops can't link
	// a session back to us if they have several nodes in the path.
	numHops := len(paymentPath)
	hopEphemeralPubKeys := make([]*btcec.PublicKey, numHops)
	hopSharedSecret := make([][sha256.Size]byte, numHops)
	hopBlindingFactors := make([][sha256.Size]byte, numHops)

	// Generate a new ephemeral key to use for ECDH for this session.
	sessionKey, err := btcec.NewPrivateKey(btcec.S256())
	if err != nil {
		return nil, nil, err
	}

	// Compute the triplet for the first hop outside of the main loop.
	// Within the loop each new triplet will be computed recursively based
	// off of the blinding factor of the last hop.
	hopEphemeralPubKeys[0] = sessionKey.PubKey()
	hopSharedSecret[0] = sha256.Sum256(btcec.GenerateSharedSecret(sessionKey, paymentPath[0]))
	hopBlindingFactors[0] = computeBlindingFactor(hopEphemeralPubKeys[0], hopSharedSecret[0][:])

	// x * b_{0} mod n
	cummulativeBlind := new(big.Int).Mul(sessionKey.X, new(big.Int).SetBytes(hopBlindingFactors[0][:]))
	cummulativeBlind.Mod(cummulativeBlind, btcec.S256().N)

	// Now recursively compute the ephemeral ECDH pub keys, the shared
	// secret, and blinding factor for each hop.
	for i := 1; i < numHops-1; i++ {
		// a_{n} = a_{n-1} x c_{n-1}
		// hopEphemeralPubKeys[i] = ScalarBaseMult(cummulativeBlind.Bytes())
		hopEphemeralPubKeys[i] = blindGroupElement(hopEphemeralPubKeys[i-1], cummulativeBlind.Bytes())

		// s_{n} = sha256( y_{n} x c_{n-1} )
		hopSharedSecret[i] = sha256.Sum256(blindGroupElement(paymentPath[i], cummulativeBlind.Bytes()).X.Bytes())

		// b_{n} = sha256(a_{n} || s_{n})
		hopBlindingFactors[i] = computeBlindingFactor(hopEphemeralPubKeys[i], hopSharedSecret[i][:])

		// c_{n} = c_{n-1} * b_{n} mod curve_order
		cummulativeBlind.Mul(cummulativeBlind, new(big.Int).SetBytes(hopBlindingFactors[i][:]))
		cummulativeBlind.Mod(cummulativeBlind, btcec.S256().N)
	}

	return nil, nil, nil
}

// ComputeBlindingFactor for the next hop given the ephemeral pubKey and
// sharedSecret for this hop. The blinding factor is computed as the
// sha-256(pubkey || sharedSecret).
func computeBlindingFactor(hopPubKey *btcec.PublicKey, hopSharedSecret []byte) [sha256.Size]byte {
	sha := sha256.New()
	sha.Write(hopPubKey.SerializeCompressed())
	sha.Write(hopSharedSecret)

	var hash [sha256.Size]byte
	copy(hash[:], sha.Sum(nil))
	return hash
}

// blindGroupElement blinds the group element by performing scalar
// multiplication of the group element by blindingFactor: G x blindingFactor.
func blindGroupElement(hopPubKey *btcec.PublicKey, blindingFactor []byte) *btcec.PublicKey {
	newX, newY := hopPubKey.Curve.ScalarMult(hopPubKey.X, hopPubKey.Y, blindingFactor[:])
	return &btcec.PublicKey{hopPubKey.Curve, newX, newY}
}

// SphinxPayload...
type SphinxPayload struct {
}

// SphinxPacket...
type SphinxPacket struct {
}
<commit_msg>Yep, secret derivation checks it out sage math<commit_after>package main

import (
	"crypto/sha256"
	"math/big"

	"github.com/btcsuite/btcd/btcec"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
)

type LnEndpoint string
type LnAddr btcutil.Address
type NodePubKey *btcec.PublicKey

type SharedSecret wire.ShaHash

var zeroNode wire.ShaHash

type MixHeader []byte

// SphinxHeader...
type SphinxHeader struct {
}

// GenerateSphinxHeader...
// TODO(roasbeef): or pass in identifiers as payment path? have map from id -> pubkey
func GenerateSphinxHeader(dest LnAddr, identifier wire.ShaHash, paymentPath []NodePubKey) (*MixHeader, []SharedSecret, error) {
	// Each hop performs ECDH with our ephemeral key pair to arrive at a
	// shared secret. Additionally, each hop randomizes the group element
	// for the next hop by multiplying it by the blinding factor. This way
	// we only need to transmit a single group element, and hops can't link
	// a session back to us if they have several nodes in the path.
	numHops := len(paymentPath)
	hopEphemeralPubKeys := make([]*btcec.PublicKey, numHops)
	hopSharedSecret := make([][sha256.Size]byte, numHops)
	hopBlindingFactors := make([][sha256.Size]byte, numHops)

	// Generate a new ephemeral key to use for ECDH for this session.
	sessionKey, err := btcec.NewPrivateKey(btcec.S256())
	if err != nil {
		return nil, nil, err
	}

	// Compute the triplet for the first hop outside of the main loop.
	// Within the loop each new triplet will be computed recursively based
	// off of the blinding factor of the last hop.
	hopEphemeralPubKeys[0] = sessionKey.PubKey()
	hopSharedSecret[0] = sha256.Sum256(btcec.GenerateSharedSecret(sessionKey, paymentPath[0]))
	hopBlindingFactors[0] = computeBlindingFactor(hopEphemeralPubKeys[0], hopSharedSecret[0][:])

	// x * b_{0} mod n. Becomes x * b_{0} * b_{1} * ..... * b_{n} mod curve_order, etc.
	cummulativeBlind := new(big.Int).Mul(sessionKey.X, new(big.Int).SetBytes(hopBlindingFactors[0][:]))
	cummulativeBlind.Mod(cummulativeBlind, btcec.S256().N)

	// Now recursively compute the ephemeral ECDH pub keys, the shared
	// secret, and blinding factor for each hop.
	for i := 1; i < numHops-1; i++ {
		// a_{n} = a_{n-1} x c_{n-1} -> (Y_prev_pub_key x prevBlindingFactor)
		hopEphemeralPubKeys[i] = blindGroupElement(hopEphemeralPubKeys[i-1], hopBlindingFactors[i-1][:])

		// s_{n} = sha256( y_{n} x c_{n-1} ) -> Y_their_pub_key x (x_our_priv * all prev blinding factors mod curve_order)
		hopSharedSecret[i] = sha256.Sum256(blindGroupElement(paymentPath[i], cummulativeBlind.Bytes()).X.Bytes())

		// b_{n} = sha256(a_{n} || s_{n})
		hopBlindingFactors[i] = computeBlindingFactor(hopEphemeralPubKeys[i], hopSharedSecret[i][:])
		// TODO(roasbeef): prob don't need to store all blinding factors, only the prev...

		// c_{n} = c_{n-1} * b_{n} mod curve_order
		cummulativeBlind.Mul(cummulativeBlind, new(big.Int).SetBytes(hopBlindingFactors[i][:]))
		cummulativeBlind.Mod(cummulativeBlind, btcec.S256().N)
	}

	return nil, nil, nil
}

// ComputeBlindingFactor for the next hop given the ephemeral pubKey and
// sharedSecret for this hop. The blinding factor is computed as the
// sha-256(pubkey || sharedSecret).
func computeBlindingFactor(hopPubKey *btcec.PublicKey, hopSharedSecret []byte) [sha256.Size]byte {
	sha := sha256.New()
	sha.Write(hopPubKey.SerializeCompressed())
	sha.Write(hopSharedSecret)

	var hash [sha256.Size]byte
	copy(hash[:], sha.Sum(nil))
	return hash
}

// blindGroupElement blinds the group element by performing scalar
// multiplication of the group element by blindingFactor: G x blindingFactor.
func blindGroupElement(hopPubKey *btcec.PublicKey, blindingFactor []byte) *btcec.PublicKey {
	newX, newY := hopPubKey.Curve.ScalarMult(hopPubKey.X, hopPubKey.Y, blindingFactor[:])
	return &btcec.PublicKey{hopPubKey.Curve, newX, newY}
}

// SphinxPayload...
type SphinxPayload struct {
}

// SphinxPacket...
type SphinxPacket struct {
}
<|endoftext|>


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

import (
	"github.com/gorilla/websocket"
	//"github.com/liangdas/mqant/log"
	"bytes"
	"io"
	"net"
	"sync"
	"time"
)

type WebsocketConnSet map[*websocket.Conn]struct{}

type WSConn struct {
	io.Reader //Read(p []byte) (n int, err error)
	io.Writer //Write(p []byte) (n int, err error)
	sync.Mutex
	buf_lock  chan error //当有写入一次数据设置一次
	buffer    bytes.Buffer
	conn      *websocket.Conn
	readfirst bool
	closeFlag bool
}

func newWSConn(conn *websocket.Conn) *WSConn {
	wsConn := new(WSConn)
	wsConn.conn = conn
	wsConn.buf_lock = make(chan error)
	wsConn.readfirst = false
	go func() {
		for {
			_, b, err := wsConn.conn.ReadMessage()
			if err != nil {
				//log.Error("读取数据失败 %s",err.Error())
				wsConn.buf_lock <- err
				break
			} else {
				wsConn.buffer.Write(b)
				wsConn.readfirst = true
				wsConn.buf_lock <- nil
			}
		}

		conn.Close()
		wsConn.Lock()
		wsConn.closeFlag = true
		wsConn.Unlock()
	}()

	return wsConn
}

func (wsConn *WSConn) doDestroy() {
	wsConn.conn.UnderlyingConn().(*net.TCPConn).SetLinger(0)
	wsConn.conn.Close()

	if !wsConn.closeFlag {
		wsConn.closeFlag = true
	}
}

func (wsConn *WSConn) Destroy() {
	wsConn.Lock()
	defer wsConn.Unlock()

	wsConn.doDestroy()
}

func (wsConn *WSConn) Close() error {
	wsConn.Lock()
	defer wsConn.Unlock()
	if wsConn.closeFlag {
		return nil
	}
	wsConn.closeFlag = true
	return wsConn.conn.Close()
}

func (wsConn *WSConn) Write(p []byte) (n int, err error) {
	err = wsConn.conn.WriteMessage(websocket.BinaryMessage, p)
	if err != nil {
		return 0, err
	}
	return len(p), nil
}

// goroutine not safe
func (wsConn *WSConn) Read(p []byte) (n int, err error) {
	err = <-wsConn.buf_lock //等待写入数据
	if err != nil {
		//读取数据出现异常了
		return
	}
	if wsConn.buffer.Len() == 0 {
		//再等一次
		err = <-wsConn.buf_lock //等待写入数据
		if err != nil {
			//读取数据出现异常了
			return
		}
	}
	return wsConn.buffer.Read(p)
}

func (wsConn *WSConn) LocalAddr() net.Addr {
	return wsConn.conn.LocalAddr()
}

func (wsConn *WSConn) RemoteAddr() net.Addr {
	return wsConn.conn.RemoteAddr()
}

// A zero value for t means I/O operations will not time out.
func (wsConn *WSConn) SetDeadline(t time.Time) error {
	err := wsConn.conn.SetWriteDeadline(t)
	if err != nil {
		return err
	}
	return wsConn.conn.SetWriteDeadline(t)
}

// SetReadDeadline sets the deadline for future Read calls.
// A zero value for t means Read will not time out.
func (wsConn *WSConn) SetReadDeadline(t time.Time) error {
	return wsConn.conn.SetReadDeadline(t)
}

// SetWriteDeadline sets the deadline for future Write calls.
// Even if write times out, it may return n > 0, indicating that
// some of the data was successfully written.
// A zero value for t means Write will not time out.
func (wsConn *WSConn) SetWriteDeadline(t time.Time) error {
	return wsConn.conn.SetWriteDeadline(t)
}
<commit_msg>websocket close<commit_after>// Copyright 2014 mqant Author. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package network

import (
	"github.com/gorilla/websocket"
	//"github.com/liangdas/mqant/log"
	"bytes"
	"io"
	"net"
	"sync"
	"time"
)

type WebsocketConnSet map[*websocket.Conn]struct{}

type WSConn struct {
	io.Reader //Read(p []byte) (n int, err error)
	io.Writer //Write(p []byte) (n int, err error)
	sync.Mutex
	buf_lock  chan error //当有写入一次数据设置一次
	buffer    bytes.Buffer
	conn      *websocket.Conn
	readfirst bool
	closeFlag bool
}

func newWSConn(conn *websocket.Conn) *WSConn {
	wsConn := new(WSConn)
	wsConn.conn = conn
	wsConn.buf_lock = make(chan error)
	wsConn.readfirst = false
	go func() {
		for {
			_, b, err := wsConn.conn.ReadMessage()
			if err != nil {
				//log.Error("读取数据失败 %s",err.Error())
				wsConn.buf_lock <- err
				break
			} else {
				wsConn.buffer.Write(b)
				wsConn.readfirst = true
				wsConn.buf_lock <- nil
			}
		}
		conn.Close()
		wsConn.Lock()
		wsConn.closeFlag = true
		wsConn.Unlock()
	}()

	return wsConn
}

func (wsConn *WSConn) doDestroy() {
	wsConn.conn.UnderlyingConn().(*net.TCPConn).SetLinger(0)
	wsConn.conn.Close()

	if !wsConn.closeFlag {
		wsConn.closeFlag = true
	}
}

func (wsConn *WSConn) Destroy() {
	wsConn.Lock()
	defer wsConn.Unlock()

	wsConn.doDestroy()
}

func (wsConn *WSConn) Close() error {
	wsConn.Lock()
	defer wsConn.Unlock()
	if wsConn.closeFlag {
		return nil
	}
	wsConn.closeFlag = true
	return wsConn.conn.Close()
}

func (wsConn *WSConn) Write(p []byte) (n int, err error) {
	err = wsConn.conn.WriteMessage(websocket.BinaryMessage, p)
	if err != nil {
		return 0, err
	}
	return len(p), nil
}

// goroutine not safe
func (wsConn *WSConn) Read(p []byte) (n int, err error) {
	err = <-wsConn.buf_lock //等待写入数据
	if err != nil {
		//读取数据出现异常了
		return
	}
	if wsConn.buffer.Len() == 0 {
		//再等一次
		err = <-wsConn.buf_lock //等待写入数据
		if err != nil {
			//读取数据出现异常了
			return
		}
	}
	return wsConn.buffer.Read(p)
}

func (wsConn *WSConn) LocalAddr() net.Addr {
	return wsConn.conn.LocalAddr()
}

func (wsConn *WSConn) RemoteAddr() net.Addr {
	return wsConn.conn.RemoteAddr()
}

// A zero value for t means I/O operations will not time out.
func (wsConn *WSConn) SetDeadline(t time.Time) error {
	err := wsConn.conn.SetWriteDeadline(t)
	if err != nil {
		return err
	}
	return wsConn.conn.SetWriteDeadline(t)
}

// SetReadDeadline sets the deadline for future Read calls.
// A zero value for t means Read will not time out.
func (wsConn *WSConn) SetReadDeadline(t time.Time) error {
	return wsConn.conn.SetReadDeadline(t)
}

// SetWriteDeadline sets the deadline for future Write calls.
// Even if write times out, it may return n > 0, indicating that
// some of the data was successfully written.
// A zero value for t means Write will not time out.
func (wsConn *WSConn) SetWriteDeadline(t time.Time) error {
	return wsConn.conn.SetWriteDeadline(t)
}
<|endoftext|>