<commit_before>// Copyright (C) 2014 Space Monkey, Inc.
//
// 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.

// +build cgo

package openssl

// #include <openssl/conf.h>
// #include <openssl/ssl.h>
// #include <openssl/x509v3.h>
//
// void OPENSSL_free_not_a_macro(void *ref) { OPENSSL_free(ref); }
//
import "C"

import (
	"errors"
	"io/ioutil"
	"runtime"
	"time"
	"unsafe"
)

type EVP_MD int

const (
	EVP_NULL      EVP_MD = iota
	EVP_MD5       EVP_MD = iota
	EVP_SHA       EVP_MD = iota
	EVP_SHA1      EVP_MD = iota
	EVP_DSS       EVP_MD = iota
	EVP_DSS1      EVP_MD = iota
	EVP_MDC2      EVP_MD = iota
	EVP_RIPEMD160 EVP_MD = iota
	EVP_SHA224    EVP_MD = iota
	EVP_SHA256    EVP_MD = iota
	EVP_SHA384    EVP_MD = iota
	EVP_SHA512    EVP_MD = iota
)

type Certificate struct {
	x      *C.X509
	Issuer *Certificate
	ref    interface{}
	pubKey PublicKey
}

type CertificateInfo struct {
	Serial       int
	Issued       time.Duration
	Expires      time.Duration
	Country      string
	Organization string
	CommonName   string
}

type Name struct {
	name *C.X509_NAME
}

// Allocate and return a new Name object.
func NewName() (*Name, error) {
	n := C.X509_NAME_new()
	if n == nil {
		return nil, errors.New("could not create x509 name")
	}
	name := &Name{name: n}
	runtime.SetFinalizer(name, func(n *Name) {
		C.X509_NAME_free(n.name)
	})
	return name, nil
}

// AddTextEntry appends a text entry to an X509 NAME.
func (n *Name) AddTextEntry(field, value string) error {
	cfield := C.CString(field)
	defer C.free(unsafe.Pointer(cfield))
	cvalue := (*C.uchar)(unsafe.Pointer(C.CString(value)))
	defer C.free(unsafe.Pointer(cvalue))
	ret := C.X509_NAME_add_entry_by_txt(
		n.name, cfield, C.MBSTRING_ASC, cvalue, -1, -1, 0)
	if ret != 1 {
		return errors.New("failed to add x509 name text entry")
	}
	return nil
}

// AddTextEntries allows adding multiple entries to a name in one call.
func (n *Name) AddTextEntries(entries map[string]string) error {
	for f, v := range entries {
		if err := n.AddTextEntry(f, v); err != nil {
			return err
		}
	}
	return nil
}

// NewCertificate generates a basic certificate based
// on the provided CertificateInfo struct
func NewCertificate(info *CertificateInfo, key PublicKey) (*Certificate, error) {
	c := &Certificate{x: C.X509_new()}
	runtime.SetFinalizer(c, func(c *Certificate) {
		C.X509_free(c.x)
	})

	name, err := c.GetSubjectName()
	if err != nil {
		return nil, err
	}
	err = name.AddTextEntries(map[string]string{
		"C":  info.Country,
		"O":  info.Organization,
		"CN": info.CommonName,
	})
	if err != nil {
		return nil, err
	}
	// self-issue for now
	if err := c.SetIssuerName(name); err != nil {
		return nil, err
	}
	if err := c.SetSerial(info.Serial); err != nil {
		return nil, err
	}
	if err := c.SetIssueDate(info.Issued); err != nil {
		return nil, err
	}
	if err := c.SetExpireDate(info.Expires); err != nil {
		return nil, err
	}
	if err := c.SetPubKey(key); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Certificate) GetSubjectName() (*Name, error) {
	n := C.X509_get_subject_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get subject name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) GetIssuerName() (*Name, error) {
	n := C.X509_get_issuer_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get issuer name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) SetSubjectName(name *Name) error {
	if C.X509_set_subject_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetIssuer updates the stored Issuer cert
// and the internal x509 Issuer Name of a certificate.
// The stored Issuer reference is used when adding extensions.
func (c *Certificate) SetIssuer(issuer *Certificate) error {
	name, err := issuer.GetSubjectName()
	if err != nil {
		return err
	}
	if err = c.SetIssuerName(name); err != nil {
		return err
	}
	c.Issuer = issuer
	return nil
}

// SetIssuerName populates the issuer name of a certificate.
// Use SetIssuer instead, if possible.
func (c *Certificate) SetIssuerName(name *Name) error {
	if C.X509_set_issuer_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetSerial sets the serial of a certificate.
func (c *Certificate) SetSerial(serial int) error {
	if C.ASN1_INTEGER_set(C.X509_get_serialNumber(c.x), C.long(serial)) != 1 {
		return errors.New("failed to set serial")
	}
	return nil
}

// SetIssueDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetIssueDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notBefore, offset)
	if result == nil {
		return errors.New("failed to set issue date")
	}
	return nil
}

// SetExpireDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetExpireDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notAfter, offset)
	if result == nil {
		return errors.New("failed to set expire date")
	}
	return nil
}

// SetPubKey assigns a new public key to a certificate.
func (c *Certificate) SetPubKey(pubKey PublicKey) error {
	c.pubKey = pubKey
	if C.X509_set_pubkey(c.x, pubKey.evpPKey()) != 1 {
		return errors.New("failed to set public key")
	}
	return nil
}

// Sign a certificate using a private key and a digest name.
// Accepted digest names are 'sha256', 'sha384', and 'sha512'.
func (c *Certificate) Sign(privKey PrivateKey, digest EVP_MD) error {
	switch digest {
	case EVP_SHA256:
	case EVP_SHA384:
	case EVP_SHA512:
	default:
		return errors.New("Unsupported digest" +
			"You're probably looking for 'EVP_SHA256' or 'EVP_SHA512'.")
	}
	return c.insecureSign(privKey, digest)
}

func (c *Certificate) insecureSign(privKey PrivateKey, digest EVP_MD) error {
	var md *C.EVP_MD
	switch digest {
	// please don't use these digest functions
	case EVP_NULL:
		md = C.EVP_md_null()
	case EVP_MD5:
		md = C.EVP_md5()
	case EVP_SHA:
		md = C.EVP_sha()
	case EVP_SHA1:
		md = C.EVP_sha1()
	case EVP_DSS:
		md = C.EVP_dss()
	case EVP_DSS1:
		md = C.EVP_dss1()
	case EVP_RIPEMD160:
		md = C.EVP_ripemd160()
	case EVP_SHA224:
		md = C.EVP_sha224()
	// you actually want one of these
	case EVP_SHA256:
		md = C.EVP_sha256()
	case EVP_SHA384:
		md = C.EVP_sha384()
	case EVP_SHA512:
		md = C.EVP_sha512()
	}
	if C.X509_sign(c.x, privKey.evpPKey(), md) <= 0 {
		return errors.New("failed to sign certificate")
	}
	return nil
}

// Add an extension to a certificate.
// Extension constants are NID_* as found in openssl.
func (c *Certificate) AddExtension(nid NID, value string) error {
	issuer := c
	if c.Issuer != nil {
		issuer = c.Issuer
	}
	var ctx C.X509V3_CTX
	C.X509V3_set_ctx(&ctx, c.x, issuer.x, nil, nil, 0)
	ex := C.X509V3_EXT_conf_nid(nil, &ctx, C.int(nid), C.CString(value))
	if ex == nil {
		return errors.New("failed to create x509v3 extension")
	}
	defer C.X509_EXTENSION_free(ex)
	if C.X509_add_ext(c.x, ex, -1) <= 0 {
		return errors.New("failed to add x509v3 extension")
	}
	return nil
}

// Wraps AddExtension using a map of NID to text extension.
// Will return without finishing if it encounters an error.
func (c *Certificate) AddExtensions(extensions map[NID]string) error {
	for nid, value := range extensions {
		if err := c.AddExtension(nid, value); err != nil {
			return err
		}
	}
	return nil
}

// LoadCertificateFromPEM loads an X509 certificate from a PEM-encoded block.
func LoadCertificateFromPEM(pem_block []byte) (*Certificate, error) {
	if len(pem_block) == 0 {
		return nil, errors.New("empty pem block")
	}
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bio := C.BIO_new_mem_buf(unsafe.Pointer(&pem_block[0]),
		C.int(len(pem_block)))
	cert := C.PEM_read_bio_X509(bio, nil, nil, nil)
	C.BIO_free(bio)
	if cert == nil {
		return nil, errorFromErrorQueue()
	}
	x := &Certificate{x: cert}
	runtime.SetFinalizer(x, func(x *Certificate) {
		C.X509_free(x.x)
	})
	return x, nil
}

// MarshalPEM converts the X509 certificate to PEM-encoded format
func (c *Certificate) MarshalPEM() (pem_block []byte, err error) {
	bio := C.BIO_new(C.BIO_s_mem())
	if bio == nil {
		return nil, errors.New("failed to allocate memory BIO")
	}
	defer C.BIO_free(bio)
	if int(C.PEM_write_bio_X509(bio, c.x)) != 1 {
		return nil, errors.New("failed dumping certificate")
	}
	return ioutil.ReadAll(asAnyBio(bio))
}

// PublicKey returns the public key embedded in the X509 certificate.
func (c *Certificate) PublicKey() (PublicKey, error) {
	pkey := C.X509_get_pubkey(c.x)
	if pkey == nil {
		return nil, errors.New("no public key found")
	}
	key := &pKey{key: pkey}
	runtime.SetFinalizer(key, func(key *pKey) {
		C.EVP_PKEY_free(key.key)
	})
	return key, nil
}

// GetSerialNumberHex returns the certificate's serial number in hex format
func (c *Certificate) GetSerialNumberHex() (serial string) {
	asn1_i := C.X509_get_serialNumber(c.x)
	bignum := C.ASN1_INTEGER_to_BN(asn1_i, nil)
	hex := C.BN_bn2hex(bignum)
	serial = C.GoString(hex)
	C.BN_free(bignum)
	C.OPENSSL_free_not_a_macro(unsafe.Pointer(hex))
	return
}
<commit_msg>Remove code restricting certificate signing to a small set of algorithms<commit_after>// Copyright (C) 2014 Space Monkey, Inc.
//
// 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.

// +build cgo

package openssl

// #include <openssl/conf.h>
// #include <openssl/ssl.h>
// #include <openssl/x509v3.h>
//
// void OPENSSL_free_not_a_macro(void *ref) { OPENSSL_free(ref); }
//
import "C"

import (
	"errors"
	"io/ioutil"
	"runtime"
	"time"
	"unsafe"
)

type EVP_MD int

const (
	EVP_NULL      EVP_MD = iota
	EVP_MD5       EVP_MD = iota
	EVP_SHA       EVP_MD = iota
	EVP_SHA1      EVP_MD = iota
	EVP_DSS       EVP_MD = iota
	EVP_DSS1      EVP_MD = iota
	EVP_MDC2      EVP_MD = iota
	EVP_RIPEMD160 EVP_MD = iota
	EVP_SHA224    EVP_MD = iota
	EVP_SHA256    EVP_MD = iota
	EVP_SHA384    EVP_MD = iota
	EVP_SHA512    EVP_MD = iota
)

type Certificate struct {
	x      *C.X509
	Issuer *Certificate
	ref    interface{}
	pubKey PublicKey
}

type CertificateInfo struct {
	Serial       int
	Issued       time.Duration
	Expires      time.Duration
	Country      string
	Organization string
	CommonName   string
}

type Name struct {
	name *C.X509_NAME
}

// Allocate and return a new Name object.
func NewName() (*Name, error) {
	n := C.X509_NAME_new()
	if n == nil {
		return nil, errors.New("could not create x509 name")
	}
	name := &Name{name: n}
	runtime.SetFinalizer(name, func(n *Name) {
		C.X509_NAME_free(n.name)
	})
	return name, nil
}

// AddTextEntry appends a text entry to an X509 NAME.
func (n *Name) AddTextEntry(field, value string) error {
	cfield := C.CString(field)
	defer C.free(unsafe.Pointer(cfield))
	cvalue := (*C.uchar)(unsafe.Pointer(C.CString(value)))
	defer C.free(unsafe.Pointer(cvalue))
	ret := C.X509_NAME_add_entry_by_txt(
		n.name, cfield, C.MBSTRING_ASC, cvalue, -1, -1, 0)
	if ret != 1 {
		return errors.New("failed to add x509 name text entry")
	}
	return nil
}

// AddTextEntries allows adding multiple entries to a name in one call.
func (n *Name) AddTextEntries(entries map[string]string) error {
	for f, v := range entries {
		if err := n.AddTextEntry(f, v); err != nil {
			return err
		}
	}
	return nil
}

// NewCertificate generates a basic certificate based
// on the provided CertificateInfo struct
func NewCertificate(info *CertificateInfo, key PublicKey) (*Certificate, error) {
	c := &Certificate{x: C.X509_new()}
	runtime.SetFinalizer(c, func(c *Certificate) {
		C.X509_free(c.x)
	})

	name, err := c.GetSubjectName()
	if err != nil {
		return nil, err
	}
	err = name.AddTextEntries(map[string]string{
		"C":  info.Country,
		"O":  info.Organization,
		"CN": info.CommonName,
	})
	if err != nil {
		return nil, err
	}
	// self-issue for now
	if err := c.SetIssuerName(name); err != nil {
		return nil, err
	}
	if err := c.SetSerial(info.Serial); err != nil {
		return nil, err
	}
	if err := c.SetIssueDate(info.Issued); err != nil {
		return nil, err
	}
	if err := c.SetExpireDate(info.Expires); err != nil {
		return nil, err
	}
	if err := c.SetPubKey(key); err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Certificate) GetSubjectName() (*Name, error) {
	n := C.X509_get_subject_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get subject name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) GetIssuerName() (*Name, error) {
	n := C.X509_get_issuer_name(c.x)
	if n == nil {
		return nil, errors.New("failed to get issuer name")
	}
	return &Name{name: n}, nil
}

func (c *Certificate) SetSubjectName(name *Name) error {
	if C.X509_set_subject_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetIssuer updates the stored Issuer cert
// and the internal x509 Issuer Name of a certificate.
// The stored Issuer reference is used when adding extensions.
func (c *Certificate) SetIssuer(issuer *Certificate) error {
	name, err := issuer.GetSubjectName()
	if err != nil {
		return err
	}
	if err = c.SetIssuerName(name); err != nil {
		return err
	}
	c.Issuer = issuer
	return nil
}

// SetIssuerName populates the issuer name of a certificate.
// Use SetIssuer instead, if possible.
func (c *Certificate) SetIssuerName(name *Name) error {
	if C.X509_set_issuer_name(c.x, name.name) != 1 {
		return errors.New("failed to set subject name")
	}
	return nil
}

// SetSerial sets the serial of a certificate.
func (c *Certificate) SetSerial(serial int) error {
	if C.ASN1_INTEGER_set(C.X509_get_serialNumber(c.x), C.long(serial)) != 1 {
		return errors.New("failed to set serial")
	}
	return nil
}

// SetIssueDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetIssueDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notBefore, offset)
	if result == nil {
		return errors.New("failed to set issue date")
	}
	return nil
}

// SetExpireDate sets the certificate issue date relative to the current time.
func (c *Certificate) SetExpireDate(when time.Duration) error {
	offset := C.long(when / time.Second)
	result := C.X509_gmtime_adj(c.x.cert_info.validity.notAfter, offset)
	if result == nil {
		return errors.New("failed to set expire date")
	}
	return nil
}

// SetPubKey assigns a new public key to a certificate.
func (c *Certificate) SetPubKey(pubKey PublicKey) error {
	c.pubKey = pubKey
	if C.X509_set_pubkey(c.x, pubKey.evpPKey()) != 1 {
		return errors.New("failed to set public key")
	}
	return nil
}

// Sign a certificate using a private key and a digest name.
// Accepted digest names are 'sha256', 'sha384', and 'sha512'.
func (c *Certificate) Sign(privKey PrivateKey, digest EVP_MD) error {
	var md *C.EVP_MD
	switch digest {
	// please don't use these digest functions
	case EVP_NULL:
		md = C.EVP_md_null()
	case EVP_MD5:
		md = C.EVP_md5()
	case EVP_SHA:
		md = C.EVP_sha()
	case EVP_SHA1:
		md = C.EVP_sha1()
	case EVP_DSS:
		md = C.EVP_dss()
	case EVP_DSS1:
		md = C.EVP_dss1()
	case EVP_RIPEMD160:
		md = C.EVP_ripemd160()
	case EVP_SHA224:
		md = C.EVP_sha224()
	// you actually want one of these
	case EVP_SHA256:
		md = C.EVP_sha256()
	case EVP_SHA384:
		md = C.EVP_sha384()
	case EVP_SHA512:
		md = C.EVP_sha512()
	}
	if C.X509_sign(c.x, privKey.evpPKey(), md) <= 0 {
		return errors.New("failed to sign certificate")
	}
	return nil
}

// Add an extension to a certificate.
// Extension constants are NID_* as found in openssl.
func (c *Certificate) AddExtension(nid NID, value string) error {
	issuer := c
	if c.Issuer != nil {
		issuer = c.Issuer
	}
	var ctx C.X509V3_CTX
	C.X509V3_set_ctx(&ctx, c.x, issuer.x, nil, nil, 0)
	ex := C.X509V3_EXT_conf_nid(nil, &ctx, C.int(nid), C.CString(value))
	if ex == nil {
		return errors.New("failed to create x509v3 extension")
	}
	defer C.X509_EXTENSION_free(ex)
	if C.X509_add_ext(c.x, ex, -1) <= 0 {
		return errors.New("failed to add x509v3 extension")
	}
	return nil
}

// Wraps AddExtension using a map of NID to text extension.
// Will return without finishing if it encounters an error.
func (c *Certificate) AddExtensions(extensions map[NID]string) error {
	for nid, value := range extensions {
		if err := c.AddExtension(nid, value); err != nil {
			return err
		}
	}
	return nil
}

// LoadCertificateFromPEM loads an X509 certificate from a PEM-encoded block.
func LoadCertificateFromPEM(pem_block []byte) (*Certificate, error) {
	if len(pem_block) == 0 {
		return nil, errors.New("empty pem block")
	}
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bio := C.BIO_new_mem_buf(unsafe.Pointer(&pem_block[0]),
		C.int(len(pem_block)))
	cert := C.PEM_read_bio_X509(bio, nil, nil, nil)
	C.BIO_free(bio)
	if cert == nil {
		return nil, errorFromErrorQueue()
	}
	x := &Certificate{x: cert}
	runtime.SetFinalizer(x, func(x *Certificate) {
		C.X509_free(x.x)
	})
	return x, nil
}

// MarshalPEM converts the X509 certificate to PEM-encoded format
func (c *Certificate) MarshalPEM() (pem_block []byte, err error) {
	bio := C.BIO_new(C.BIO_s_mem())
	if bio == nil {
		return nil, errors.New("failed to allocate memory BIO")
	}
	defer C.BIO_free(bio)
	if int(C.PEM_write_bio_X509(bio, c.x)) != 1 {
		return nil, errors.New("failed dumping certificate")
	}
	return ioutil.ReadAll(asAnyBio(bio))
}

// PublicKey returns the public key embedded in the X509 certificate.
func (c *Certificate) PublicKey() (PublicKey, error) {
	pkey := C.X509_get_pubkey(c.x)
	if pkey == nil {
		return nil, errors.New("no public key found")
	}
	key := &pKey{key: pkey}
	runtime.SetFinalizer(key, func(key *pKey) {
		C.EVP_PKEY_free(key.key)
	})
	return key, nil
}

// GetSerialNumberHex returns the certificate's serial number in hex format
func (c *Certificate) GetSerialNumberHex() (serial string) {
	asn1_i := C.X509_get_serialNumber(c.x)
	bignum := C.ASN1_INTEGER_to_BN(asn1_i, nil)
	hex := C.BN_bn2hex(bignum)
	serial = C.GoString(hex)
	C.BN_free(bignum)
	C.OPENSSL_free_not_a_macro(unsafe.Pointer(hex))
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
)

// A Request is a user-configured request (which can be used to make an http.Request).
type Request struct {
	Method   string
	URL      string
	Body     string
	BodyFile string `json:"body_file"`
	Headers  map[string]string
	Weight   float64
}

type ReportingStats struct {
	Quantiles []float64
	Buckets   []float64
}

type TargetQPS struct {
	unlimited bool
	qps int
}

type Conf struct {
	Requests        []*Request
	RequestsFile    string          `json:"requests_file"`
	TargetQPS       TargetQPS       `json:"target_qps"`
	DurationSeconds int             `json:"duration_seconds"`
	MaxConcurrent   int             `json:"max_concurrent"`
	ReportingStats  *ReportingStats `json:"reporting_stats"`
}

func (q *TargetQPS) UnmarshalJSON(b []byte) error {
	if bytes.Equal(b, []byte(`"unlimited"`)) {
		q.unlimited = true
		return nil
	}
	q.unlimited = false
	if err := json.Unmarshal(b, &q.qps); err != nil {
		return err
	}
	return nil
}

func parseConfig() (*Conf, error) {
	flag.Parse()
	conf := &Conf{}
	f, err := os.Open(*confFile)
	if err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(f)
	if err := decoder.Decode(conf); err != nil {
		return nil, err
	}

	// Sanity checking
	if conf.RequestsFile != "" {
		reqFile, err := os.Open(conf.RequestsFile)
		if err != nil {
			return nil, fmt.Errorf("Cannot open requests file %s: %s", conf.RequestsFile, err)
		}
		reqDecoder := json.NewDecoder(reqFile)
		requests := []*Request{}
		if err := reqDecoder.Decode(requests); err != nil {
			return nil, fmt.Errorf("Cannot parse json requests file %s: %s", conf.RequestsFile, err)
		}
		conf.Requests = append(conf.Requests, requests...)
	}
	for _, request := range conf.Requests {
		if request.Weight == 0 {
			request.Weight = 1
		}
	}
	if len(conf.Requests) == 0 {
		return nil, fmt.Errorf("No requests supplied")
	}
	for _, r := range conf.Requests {
		if r.BodyFile != "" {
			if r.Body != "" {
				return nil, fmt.Errorf("Bad request: both body and body_file specified.")
			}
			contents, err := ioutil.ReadFile(r.BodyFile)
			if err != nil {
				return nil, fmt.Errorf("Bad request: error reading body_file %s: %s", r.BodyFile, err)
			}
			r.Body = string(contents)
		}
	}
	if conf.MaxConcurrent == 0 {
		conf.MaxConcurrent = 100
	}

	if conf.ReportingStats != nil {
		sort.Float64s(conf.ReportingStats.Quantiles)
		sort.Float64s(conf.ReportingStats.Buckets)
		for _, q := range conf.ReportingStats.Quantiles {
			if q < 0 || q >= 1 {
				return nil, fmt.Errorf("Bad quantile %f", q)
			}
		}
		if len(conf.ReportingStats.Buckets) > 0 {
			last := conf.ReportingStats.Buckets[0]
			for i, b := range conf.ReportingStats.Buckets {
				if b <= 0 {
					return nil, fmt.Errorf("Bucket is <= 0: %f", b)
				}
				if i > 0 && b == last {
					return nil, fmt.Errorf("Duplicate buckets: %f", b)
				}
				last = b
			}
		}
	}
	return conf, nil
}
<commit_msg>Fix request_file bug<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
)

// A Request is a user-configured request (which can be used to make an http.Request).
type Request struct {
	Method   string
	URL      string
	Body     string
	BodyFile string `json:"body_file"`
	Headers  map[string]string
	Weight   float64
}

type ReportingStats struct {
	Quantiles []float64
	Buckets   []float64
}

type TargetQPS struct {
	unlimited bool
	qps int
}

type Conf struct {
	Requests        []*Request
	RequestsFile    string          `json:"requests_file"`
	TargetQPS       TargetQPS       `json:"target_qps"`
	DurationSeconds int             `json:"duration_seconds"`
	MaxConcurrent   int             `json:"max_concurrent"`
	ReportingStats  *ReportingStats `json:"reporting_stats"`
}

func (q *TargetQPS) UnmarshalJSON(b []byte) error {
	if bytes.Equal(b, []byte(`"unlimited"`)) {
		q.unlimited = true
		return nil
	}
	q.unlimited = false
	if err := json.Unmarshal(b, &q.qps); err != nil {
		return err
	}
	return nil
}

func parseConfig() (*Conf, error) {
	flag.Parse()
	conf := &Conf{}
	f, err := os.Open(*confFile)
	if err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(f)
	if err := decoder.Decode(conf); err != nil {
		return nil, err
	}

	// Sanity checking
	if conf.RequestsFile != "" {
		reqFile, err := os.Open(conf.RequestsFile)
		if err != nil {
			return nil, fmt.Errorf("Cannot open requests file %s: %s", conf.RequestsFile, err)
		}
		reqDecoder := json.NewDecoder(reqFile)
		requests := []*Request{}
		if err := reqDecoder.Decode(&requests); err != nil {
			return nil, fmt.Errorf("Cannot parse json requests file %s: %s", conf.RequestsFile, err)
		}
		conf.Requests = append(conf.Requests, requests...)
	}
	for _, request := range conf.Requests {
		if request.Weight == 0 {
			request.Weight = 1
		}
	}
	if len(conf.Requests) == 0 {
		return nil, fmt.Errorf("No requests supplied")
	}
	for _, r := range conf.Requests {
		if r.BodyFile != "" {
			if r.Body != "" {
				return nil, fmt.Errorf("Bad request: both body and body_file specified.")
			}
			contents, err := ioutil.ReadFile(r.BodyFile)
			if err != nil {
				return nil, fmt.Errorf("Bad request: error reading body_file %s: %s", r.BodyFile, err)
			}
			r.Body = string(contents)
		}
	}
	if conf.MaxConcurrent == 0 {
		conf.MaxConcurrent = 100
	}

	if conf.ReportingStats != nil {
		sort.Float64s(conf.ReportingStats.Quantiles)
		sort.Float64s(conf.ReportingStats.Buckets)
		for _, q := range conf.ReportingStats.Quantiles {
			if q < 0 || q >= 1 {
				return nil, fmt.Errorf("Bad quantile %f", q)
			}
		}
		if len(conf.ReportingStats.Buckets) > 0 {
			last := conf.ReportingStats.Buckets[0]
			for i, b := range conf.ReportingStats.Buckets {
				if b <= 0 {
					return nil, fmt.Errorf("Bucket is <= 0: %f", b)
				}
				if i > 0 && b == last {
					return nil, fmt.Errorf("Duplicate buckets: %f", b)
				}
				last = b
			}
		}
	}
	return conf, nil
}
<|endoftext|>


<commit_before>package proxyprotocol

import (
	"bufio"
	"bytes"
	parser "github.com/racker/go-proxy-protocol"
	"io"
	"net"
	"time"
)

type conn struct {
	raw   net.Conn
	proxy *parser.ProxyLine
	err   error
}

func (c *conn) init() {
	if c.proxy != nil || c.err != nil {
		return
	}

	// prepare for the longest (TCP6) line
	buf := make([]byte, 107)

	// read at least the shortest (UNKNOWN) line
	n, err := c.raw.Read(buf[:15])
	if err != nil && !(err == io.EOF && n == 15) {
		c.err = err
		return
	}
	if n < 15 {
		c.err = io.EOF
		return
	}

	for i := n; i < len(buf); i++ {
		if buf[i-2] == byte('\r') {
			if buf[i-1] == byte('\n') {
				buf = buf[:i]
				break
			} else {
				c.err = parser.InvalidProxyLine
				return
			}
		}

		n, c.err = c.raw.Read(buf[i : i+1])
		if c.err != nil {
			return
		}
		if n != 1 {
			c.err = io.EOF
			return
		}
	}

	if bytes.HasPrefix(buf, []byte("PROXY UNKNOWN")) {
		c.proxy = new(parser.ProxyLine)
		return
	}

	c.proxy, c.err = parser.ConsumeProxyLine(bufio.NewReader(bytes.NewReader(buf)))
	if c.proxy == nil && c.err == nil {
		c.err = parser.InvalidProxyLine
	}
	return
}

func (c *conn) Read(b []byte) (n int, err error) {
	c.init()
	if c.err != nil {
		err = c.err
	} else {
		n, err = c.raw.Read(b)
	}
	return
}

func (c *conn) Write(b []byte) (n int, err error) {
	n, err = c.raw.Write(b)
	return
}

func (c *conn) Close() error {
	return c.raw.Close()
}

func (c *conn) LocalAddr() net.Addr {
	c.init()
	if c.err == nil && c.proxy.DstAddr != nil {
		return c.proxy.DstAddr
	}
	return c.raw.LocalAddr()
}

func (c *conn) RemoteAddr() net.Addr {
	c.init()
	if c.err == nil && c.proxy.SrcAddr != nil {
		return c.proxy.SrcAddr
	}
	return c.raw.RemoteAddr()
}

func (c *conn) SetDeadline(t time.Time) error {
	return c.raw.SetDeadline(t)
}

func (c *conn) SetReadDeadline(t time.Time) error {
	return c.raw.SetReadDeadline(t)
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	return c.raw.SetWriteDeadline(t)
}
<commit_msg>pass non-PROXY input through<commit_after>package proxyprotocol

import (
	"bufio"
	"bytes"
	parser "github.com/racker/go-proxy-protocol"
	"io"
	"net"
	"time"
)

type conn struct {
	raw   net.Conn
	buf   []byte
	err   error
	proxy *parser.ProxyLine
}

func (c *conn) init() {
	if c.buf != nil || c.err != nil || c.proxy != nil {
		return
	}

	// prepare for the longest (TCP6) line
	buf := make([]byte, 107)

	// read at least the shortest (UNKNOWN) line
	n, err := c.raw.Read(buf[:15])
	if err != nil && !(err == io.EOF && n == 15) {
		c.buf = buf[:n]
		c.err = err
		return
	}
	if n < 15 {
		c.buf = buf[:n]
		c.err = io.EOF
		return
	}

	for i := n; i < len(buf); i++ {
		if buf[i-2] == byte('\r') {
			if buf[i-1] == byte('\n') {
				buf = buf[:i]
				break
			} else {
				c.buf = buf[:i]
				return
			}
		}

		n, c.err = c.raw.Read(buf[i : i+1])
		if c.err != nil {
			c.buf = buf[:i+n]
			return
		}
		if n == 0 {
			c.buf = buf[:i]
			c.err = io.EOF
			return
		}
	}

	if bytes.HasPrefix(buf, []byte("PROXY UNKNOWN")) {
		c.proxy = new(parser.ProxyLine)
		return
	}

	c.proxy, c.err = parser.ConsumeProxyLine(bufio.NewReader(bytes.NewReader(buf)))
	if c.err == nil && c.proxy == nil {
		c.buf = buf
	}
	return
}

func (c *conn) Read(b []byte) (n int, err error) {
	c.init()

	if len(c.buf) > 0 {
		n = copy(b, c.buf)
		b = b[n:]
		c.buf = c.buf[:n]

		if len(b) == 0 {
			return
		}
	}

	if c.err != nil {
		err = c.err
		return
	}

	var l int

	l, err = c.raw.Read(b)
	n += l

	return
}

func (c *conn) Write(b []byte) (n int, err error) {
	n, err = c.raw.Write(b)
	return
}

func (c *conn) Close() error {
	c.buf = nil
	return c.raw.Close()
}

func (c *conn) LocalAddr() net.Addr {
	c.init()

	if c.proxy != nil && c.proxy.DstAddr != nil {
		return c.proxy.DstAddr
	}

	return c.raw.LocalAddr()
}

func (c *conn) RemoteAddr() net.Addr {
	c.init()

	if c.proxy != nil && c.proxy.SrcAddr != nil {
		return c.proxy.SrcAddr
	}

	return c.raw.RemoteAddr()
}

func (c *conn) SetDeadline(t time.Time) error {
	return c.raw.SetDeadline(t)
}

func (c *conn) SetReadDeadline(t time.Time) error {
	return c.raw.SetReadDeadline(t)
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	return c.raw.SetWriteDeadline(t)
}
<|endoftext|>


<commit_before>package utp

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

type Conn struct {
	s          *C.utp_socket
	cond       sync.Cond
	readBuf    []byte
	gotEOF     bool
	gotConnect bool
	// Set on state changed to UTP_STATE_DESTROYING. Not valid to refer to the
	// socket after getting this.
	destroyed bool
	// Conn.Close was called.
	closed bool

	writeDeadline time.Time
	readDeadline  time.Time
}

func (c *Conn) setConnected() {
	c.gotConnect = true
	c.cond.Broadcast()
}

func (c *Conn) waitForConnect() error {
	for !c.gotConnect {
		c.cond.Wait()
	}
	return nil
}

func (c *Conn) Close() (err error) {
	mu.Lock()
	defer mu.Unlock()
	c.close()
	return nil
}

func (c *Conn) close() {
	if c.closed {
		return
	}
	if !c.destroyed {
		C.utp_close(c.s)
		c.s = nil
	}
	c.closed = true
	c.cond.Broadcast()
}

func (c *Conn) LocalAddr() net.Addr {
	mu.Lock()
	defer mu.Unlock()
	return getSocketForLibContext(C.utp_get_context(c.s)).pc.LocalAddr()
}

func (c *Conn) Read(b []byte) (int, error) {
	mu.Lock()
	defer mu.Unlock()
	for {
		n := copy(b, c.readBuf)
		c.readBuf = c.readBuf[n:]
		log.Println(n, len(c.readBuf))
		if n != 0 && len(c.readBuf) == 0 {
			C.utp_read_drained(c.s)
		}
		if len(c.readBuf) == 0 && c.gotEOF {
			return n, io.EOF
		}
		if n != 0 || len(b) == 0 {
			return n, nil
		}
		if c.destroyed {
			return 0, errors.New("destroyed")
		}
		if c.closed {
			return 0, errors.New("closed")
		}
		if !c.readDeadline.IsZero() && !time.Now().Before(c.readDeadline) {
			return 0, errDeadlineExceeded{}
		}
		c.cond.Wait()
	}
}

func (c *Conn) Write(b []byte) (int, error) {
	mu.Lock()
	defer mu.Unlock()
	n := 0
	for len(b) != 0 {
		if c.closed {
			return n, errors.New("closed")
		}
		if c.destroyed {
			return n, errors.New("destroyed")
		}
		n1 := C.utp_write(c.s, unsafe.Pointer(&b[0]), C.size_t(len(b)))
		if n1 < 0 {
			panic(n1)
		}
		b = b[n1:]
		n += int(n1)
		if n1 == 0 {
			if !c.writeDeadline.IsZero() && !time.Now().Before(c.writeDeadline) {
				return n, errDeadlineExceeded{}
			}
			c.cond.Wait()
		}
	}
	return n, nil
}

func (c *Conn) RemoteAddr() net.Addr {
	var rsa syscall.RawSockaddrAny
	var addrlen C.socklen_t = syscall.SizeofSockaddrAny
	C.utp_getpeername(c.s, (*C.struct_sockaddr)(unsafe.Pointer(&rsa)), &addrlen)
	sa, err := anyToSockaddr(&rsa)
	if err != nil {
		panic(err)
	}
	return sockaddrToUDP(sa)
}

func (c *Conn) SetDeadline(t time.Time) error {
	mu.Lock()
	defer mu.Unlock()
	c.readDeadline = t
	c.writeDeadline = t
	c.cond.Broadcast()
	return nil
}
func (c *Conn) SetReadDeadline(t time.Time) error {
	mu.Lock()
	defer mu.Unlock()
	c.readDeadline = t
	c.cond.Broadcast()
	return nil
}
func (c *Conn) SetWriteDeadline(t time.Time) error {
	mu.Lock()
	defer mu.Unlock()
	c.writeDeadline = t
	c.cond.Broadcast()
	return nil
}

func (c *Conn) setGotEOF() {
	c.gotEOF = true
	c.cond.Broadcast()
}

func (c *Conn) onDestroyed() {
	c.destroyed = true
	c.s = nil
	c.cond.Broadcast()
}
<commit_msg>Restructure Conn.Read<commit_after>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"errors"
	"io"
	"net"
	"sync"
	"syscall"
	"time"
	"unsafe"
)

type Conn struct {
	s          *C.utp_socket
	cond       sync.Cond
	readBuf    []byte
	gotEOF     bool
	gotConnect bool
	// Set on state changed to UTP_STATE_DESTROYING. Not valid to refer to the
	// socket after getting this.
	destroyed bool
	// Conn.Close was called.
	closed bool

	writeDeadline time.Time
	readDeadline  time.Time
}

func (c *Conn) setConnected() {
	c.gotConnect = true
	c.cond.Broadcast()
}

func (c *Conn) waitForConnect() error {
	for !c.gotConnect {
		c.cond.Wait()
	}
	return nil
}

func (c *Conn) Close() (err error) {
	mu.Lock()
	defer mu.Unlock()
	c.close()
	return nil
}

func (c *Conn) close() {
	if c.closed {
		return
	}
	if !c.destroyed {
		C.utp_close(c.s)
		c.s = nil
	}
	c.closed = true
	c.cond.Broadcast()
}

func (c *Conn) LocalAddr() net.Addr {
	mu.Lock()
	defer mu.Unlock()
	return getSocketForLibContext(C.utp_get_context(c.s)).pc.LocalAddr()
}

func (c *Conn) readNoWait(b []byte) (n int, err error) {
	n = copy(b, c.readBuf)
	c.readBuf = c.readBuf[n:]
	if n != 0 && len(c.readBuf) == 0 {
		C.utp_read_drained(c.s)
	}
	err = func() error {
		switch {
		case c.gotEOF:
			return io.EOF
		case c.destroyed:
			return errors.New("destroyed")
		case c.closed:
			return errors.New("closed")
		case !c.readDeadline.IsZero() && !time.Now().Before(c.readDeadline):
			return errDeadlineExceeded{}
		default:
			return nil
		}
	}()
	return
}

func (c *Conn) Read(b []byte) (int, error) {
	mu.Lock()
	defer mu.Unlock()
	for {
		n, err := c.readNoWait(b)
		if n != 0 || len(b) == 0 || err != nil {
			return n, err
		}
		c.cond.Wait()
	}
}

func (c *Conn) Write(b []byte) (int, error) {
	mu.Lock()
	defer mu.Unlock()
	n := 0
	for len(b) != 0 {
		if c.closed {
			return n, errors.New("closed")
		}
		if c.destroyed {
			return n, errors.New("destroyed")
		}
		n1 := C.utp_write(c.s, unsafe.Pointer(&b[0]), C.size_t(len(b)))
		if n1 < 0 {
			panic(n1)
		}
		b = b[n1:]
		n += int(n1)
		if n1 == 0 {
			if !c.writeDeadline.IsZero() && !time.Now().Before(c.writeDeadline) {
				return n, errDeadlineExceeded{}
			}
			c.cond.Wait()
		}
	}
	return n, nil
}

func (c *Conn) RemoteAddr() net.Addr {
	var rsa syscall.RawSockaddrAny
	var addrlen C.socklen_t = syscall.SizeofSockaddrAny
	C.utp_getpeername(c.s, (*C.struct_sockaddr)(unsafe.Pointer(&rsa)), &addrlen)
	sa, err := anyToSockaddr(&rsa)
	if err != nil {
		panic(err)
	}
	return sockaddrToUDP(sa)
}

func (c *Conn) SetDeadline(t time.Time) error {
	mu.Lock()
	defer mu.Unlock()
	c.readDeadline = t
	c.writeDeadline = t
	c.cond.Broadcast()
	return nil
}
func (c *Conn) SetReadDeadline(t time.Time) error {
	mu.Lock()
	defer mu.Unlock()
	c.readDeadline = t
	c.cond.Broadcast()
	return nil
}
func (c *Conn) SetWriteDeadline(t time.Time) error {
	mu.Lock()
	defer mu.Unlock()
	c.writeDeadline = t
	c.cond.Broadcast()
	return nil
}

func (c *Conn) setGotEOF() {
	c.gotEOF = true
	c.cond.Broadcast()
}

func (c *Conn) onDestroyed() {
	c.destroyed = true
	c.s = nil
	c.cond.Broadcast()
}
<|endoftext|>


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

package websocket

import (
	"bytes"
	"fmt"
	"io"
	"net/url"
	"time"

	"github.com/gopherjs/gopherjs/js"
	"honnef.co/go/js/dom"
)

func beginHandlerOpen(ch chan error, removeHandlers func()) func(ev js.Object) {
	return func(ev js.Object) {
		removeHandlers()
		close(ch)
	}
}

// closeError allows a CloseEvent to be used as an error.
type closeError struct {
	*dom.CloseEvent
}

func (e *closeError) Error() string {
	var cleanStmt string
	if e.WasClean {
		cleanStmt = "clean"
	} else {
		cleanStmt = "unclean"
	}
	return fmt.Sprintf("CloseEvent: (%s) (%d) %s", cleanStmt, e.Code, e.Reason)
}

func beginHandlerClose(ch chan error, removeHandlers func()) func(ev js.Object) {
	return func(ev js.Object) {
		removeHandlers()
		go func() {
			ce := dom.WrapEvent(ev).(*dom.CloseEvent)
			ch <- &closeError{CloseEvent: ce}
			close(ch)
		}()
	}
}

type deadlineErr struct{}

func (e *deadlineErr) Error() string   { return "i/o timeout: deadline reached" }
func (e *deadlineErr) Timeout() bool   { return true }
func (e *deadlineErr) Temporary() bool { return true }

var errDeadlineReached = &deadlineErr{}

// TODO(nightexcessive): Add a Dial function that allows a deadline to be
// specified.

// Dial opens a new WebSocket connection. It will block until the connection is
// established or fails to connect.
func Dial(url string) (*Conn, error) {
	ws, err := New(url)
	if err != nil {
		return nil, err
	}
	conn := &Conn{
		WebSocket: ws,
		ch:        make(chan *dom.MessageEvent, 1),
	}
	conn.initialize()

	openCh := make(chan error, 1)

	var (
		openHandler  func(ev js.Object)
		closeHandler func(ev js.Object)
	)

	// Handlers need to be removed to prevent a panic when the WebSocket closes
	// immediately and fires both open and close before they can be removed.
	// This way, handlers are removed before the channel is closed.
	removeHandlers := func() {
		ws.RemoveEventListener("open", false, openHandler)
		ws.RemoveEventListener("close", false, closeHandler)
	}

	// We have to use variables for the functions so that we can remove the
	// event handlers afterwards.
	openHandler = beginHandlerOpen(openCh, removeHandlers)
	closeHandler = beginHandlerClose(openCh, removeHandlers)

	ws.AddEventListener("open", false, openHandler)
	ws.AddEventListener("close", false, closeHandler)

	err, ok := <-openCh
	if ok && err != nil {
		return nil, err
	}

	return conn, nil
}

// Conn is a high-level wrapper around WebSocket. It is intended to satisfy the
// net.Conn interface.
//
// To create a Conn, use Dial. Instantiating Conn without Dial will not work.
type Conn struct {
	*WebSocket

	ch      chan *dom.MessageEvent
	readBuf *bytes.Reader

	readDeadline time.Time
}

func (c *Conn) onMessage(event js.Object) {
	go func() {
		c.ch <- dom.WrapEvent(event).(*dom.MessageEvent)
	}()
}

func (c *Conn) onClose(event js.Object) {
	go func() {
		// We queue nil to the end so that any messages received prior to
		// closing get handled first.
		c.ch <- nil
	}()
}

// initialize adds all of the event handlers necessary for a Conn to function.
// It should never be called more than once and is already called if Dial was
// used to create the Conn.
func (c *Conn) initialize() {
	// We need this so that received binary data is in ArrayBufferView format so
	// that it can easily be read.
	c.BinaryType = "arraybuffer"

	c.AddEventListener("message", false, c.onMessage)
	c.AddEventListener("close", false, c.onClose)
}

// handleFrame handles a single frame received from the channel. This is a
// convenience funciton to dedupe code for the multiple deadline cases.
func (c *Conn) handleFrame(item *dom.MessageEvent, ok bool) (*dom.MessageEvent, error) {
	if !ok { // The channel has been closed
		return nil, io.EOF
	} else if item == nil {
		// See onClose for the explanation about sending a nil item.
		close(c.ch)
		return nil, io.EOF
	}

	return item, nil
}

// receiveFrame receives one full frame from the WebSocket. It blocks until the
// frame is received.
func (c *Conn) receiveFrame(observeDeadline bool) (*dom.MessageEvent, error) {
	var deadlineChan <-chan time.Time // Receiving on a nil channel always blocks indefinitely

	if observeDeadline && !c.readDeadline.IsZero() {
		now := time.Now()
		if now.After(c.readDeadline) {
			select {
			case item, ok := <-c.ch:
				return c.handleFrame(item, ok)
			default:
				return nil, errDeadlineReached
			}
		}

		timer := time.NewTimer(c.readDeadline.Sub(now))
		defer timer.Stop()

		deadlineChan = timer.C
	}

	select {
	case item, ok := <-c.ch:
		return c.handleFrame(item, ok)
	case <-deadlineChan:
		return nil, errDeadlineReached
	}
}

func getFrameData(obj js.Object) []byte {
	// Check if it's an array buffer. If so, convert it to a Go byte slice.
	if constructor := obj.Get("constructor"); constructor == js.Global.Get("ArrayBuffer") {
		int8Array := js.Global.Get("Uint8Array").New(obj)
		return int8Array.Interface().([]byte)
	}

	return []byte(obj.String())
}

func (c *Conn) Read(b []byte) (n int, err error) {
	if c.readBuf != nil {
		n, err = c.readBuf.Read(b)
		if err == io.EOF {
			c.readBuf = nil
			err = nil
		}
		// If we read nothing from the buffer, continue to trying to receive.
		// This saves us when the last Read call emptied the buffer and this
		// call triggers the EOF. There's probably a better way of doing this,
		// but I'm really tired.
		if n > 0 {
			return
		}
	}

	frame, err := c.receiveFrame(true)
	if err != nil {
		return 0, err
	}

	receivedBytes := getFrameData(frame.Data)

	n = copy(b, receivedBytes)
	// Fast path: The entire frame's contents have been copied into b.
	if n >= len(receivedBytes) {
		return
	}

	c.readBuf = bytes.NewReader(receivedBytes[n:])
	return
}

// Write writes the contents of b to the WebSocket using a binary opcode.
func (c *Conn) Write(b []byte) (n int, err error) {
	// []byte is converted to an (U)Int8Array by GopherJS, which fullfils the
	// ArrayBufferView definition.
	err = c.Send(b)
	if err != nil {
		return
	}
	n = len(b)
	return
}

// WriteString writes the contents of s to the WebSocket using a text frame
// opcode.
func (c *Conn) WriteString(s string) (n int, err error) {
	err = c.Send(s)
	if err != nil {
		return
	}
	n = len(s)
	return
}

// BUG(nightexcessive): Conn doesn't currently fulfill the net.Conn interface
// because we can't return net.Addr from Conn.LocalAddr and Conn.RemoteAddr
// because net.init() causes a panic due to attempts to make syscalls.
//
// See: https://github.com/gopherjs/gopherjs/issues/123

// LocalAddr would typically return the local network address, but due to
// limitations in the JavaScript API, it is unable to. Calling this method will
// cause a panic.
func (c *Conn) LocalAddr() *Addr {
	// BUG(nightexcessive): Conn.LocalAddr() panics because the underlying
	// JavaScript API has no way of figuring out the local address.

	// TODO(nightexcessive): Find a more graceful way to handle this
	panic("we are unable to implement websocket.Conn.LocalAddr() due to limitations in the underlying JavaScript API")
}

// RemoteAddr returns the remote network address, based on
// websocket.WebSocket.URL.
func (c *Conn) RemoteAddr() *Addr {
	wsURL, err := url.Parse(c.URL)
	if err != nil {
		// TODO(nightexcessive): Should we be panicking for this?
		panic(err)
	}
	return &Addr{wsURL}
}

// SetDeadline sets the read and write deadlines associated with the connection.
// It is equivalent to calling both SetReadDeadline and SetWriteDeadline.
//
// A zero value for t means that I/O operations will not time out.
func (c *Conn) SetDeadline(t time.Time) error {
	c.readDeadline = t
	return nil
}

// SetReadDeadline sets the deadline for future Read calls. A zero value for t
// means Read will not time out.
func (c *Conn) SetReadDeadline(t time.Time) error {
	c.readDeadline = t
	return nil
}

// SetWriteDeadline sets the deadline for future Write calls. Because our writes
// do not block, this function is a no-op.
func (c *Conn) SetWriteDeadline(t time.Time) error {
	return nil
}
<commit_msg>Fulfill net.Conn interface<commit_after>// Copyright 2014-2015 GopherJS Team. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file.

package websocket

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/url"
	"time"

	"github.com/gopherjs/gopherjs/js"
	"honnef.co/go/js/dom"
)

func beginHandlerOpen(ch chan error, removeHandlers func()) func(ev js.Object) {
	return func(ev js.Object) {
		removeHandlers()
		close(ch)
	}
}

// closeError allows a CloseEvent to be used as an error.
type closeError struct {
	*dom.CloseEvent
}

func (e *closeError) Error() string {
	var cleanStmt string
	if e.WasClean {
		cleanStmt = "clean"
	} else {
		cleanStmt = "unclean"
	}
	return fmt.Sprintf("CloseEvent: (%s) (%d) %s", cleanStmt, e.Code, e.Reason)
}

func beginHandlerClose(ch chan error, removeHandlers func()) func(ev js.Object) {
	return func(ev js.Object) {
		removeHandlers()
		go func() {
			ce := dom.WrapEvent(ev).(*dom.CloseEvent)
			ch <- &closeError{CloseEvent: ce}
			close(ch)
		}()
	}
}

type deadlineErr struct{}

func (e *deadlineErr) Error() string   { return "i/o timeout: deadline reached" }
func (e *deadlineErr) Timeout() bool   { return true }
func (e *deadlineErr) Temporary() bool { return true }

var errDeadlineReached = &deadlineErr{}

// TODO(nightexcessive): Add a Dial function that allows a deadline to be
// specified.

// Dial opens a new WebSocket connection. It will block until the connection is
// established or fails to connect.
func Dial(url string) (*Conn, error) {
	ws, err := New(url)
	if err != nil {
		return nil, err
	}
	conn := &Conn{
		WebSocket: ws,
		ch:        make(chan *dom.MessageEvent, 1),
	}
	conn.initialize()

	openCh := make(chan error, 1)

	var (
		openHandler  func(ev js.Object)
		closeHandler func(ev js.Object)
	)

	// Handlers need to be removed to prevent a panic when the WebSocket closes
	// immediately and fires both open and close before they can be removed.
	// This way, handlers are removed before the channel is closed.
	removeHandlers := func() {
		ws.RemoveEventListener("open", false, openHandler)
		ws.RemoveEventListener("close", false, closeHandler)
	}

	// We have to use variables for the functions so that we can remove the
	// event handlers afterwards.
	openHandler = beginHandlerOpen(openCh, removeHandlers)
	closeHandler = beginHandlerClose(openCh, removeHandlers)

	ws.AddEventListener("open", false, openHandler)
	ws.AddEventListener("close", false, closeHandler)

	err, ok := <-openCh
	if ok && err != nil {
		return nil, err
	}

	return conn, nil
}

// Conn is a high-level wrapper around WebSocket. It is intended to satisfy the
// net.Conn interface.
//
// To create a Conn, use Dial. Instantiating Conn without Dial will not work.
type Conn struct {
	*WebSocket

	ch      chan *dom.MessageEvent
	readBuf *bytes.Reader

	readDeadline time.Time
}

func (c *Conn) onMessage(event js.Object) {
	go func() {
		c.ch <- dom.WrapEvent(event).(*dom.MessageEvent)
	}()
}

func (c *Conn) onClose(event js.Object) {
	go func() {
		// We queue nil to the end so that any messages received prior to
		// closing get handled first.
		c.ch <- nil
	}()
}

// initialize adds all of the event handlers necessary for a Conn to function.
// It should never be called more than once and is already called if Dial was
// used to create the Conn.
func (c *Conn) initialize() {
	// We need this so that received binary data is in ArrayBufferView format so
	// that it can easily be read.
	c.BinaryType = "arraybuffer"

	c.AddEventListener("message", false, c.onMessage)
	c.AddEventListener("close", false, c.onClose)
}

// handleFrame handles a single frame received from the channel. This is a
// convenience funciton to dedupe code for the multiple deadline cases.
func (c *Conn) handleFrame(item *dom.MessageEvent, ok bool) (*dom.MessageEvent, error) {
	if !ok { // The channel has been closed
		return nil, io.EOF
	} else if item == nil {
		// See onClose for the explanation about sending a nil item.
		close(c.ch)
		return nil, io.EOF
	}

	return item, nil
}

// receiveFrame receives one full frame from the WebSocket. It blocks until the
// frame is received.
func (c *Conn) receiveFrame(observeDeadline bool) (*dom.MessageEvent, error) {
	var deadlineChan <-chan time.Time // Receiving on a nil channel always blocks indefinitely

	if observeDeadline && !c.readDeadline.IsZero() {
		now := time.Now()
		if now.After(c.readDeadline) {
			select {
			case item, ok := <-c.ch:
				return c.handleFrame(item, ok)
			default:
				return nil, errDeadlineReached
			}
		}

		timer := time.NewTimer(c.readDeadline.Sub(now))
		defer timer.Stop()

		deadlineChan = timer.C
	}

	select {
	case item, ok := <-c.ch:
		return c.handleFrame(item, ok)
	case <-deadlineChan:
		return nil, errDeadlineReached
	}
}

func getFrameData(obj js.Object) []byte {
	// Check if it's an array buffer. If so, convert it to a Go byte slice.
	if constructor := obj.Get("constructor"); constructor == js.Global.Get("ArrayBuffer") {
		int8Array := js.Global.Get("Uint8Array").New(obj)
		return int8Array.Interface().([]byte)
	}

	return []byte(obj.String())
}

func (c *Conn) Read(b []byte) (n int, err error) {
	if c.readBuf != nil {
		n, err = c.readBuf.Read(b)
		if err == io.EOF {
			c.readBuf = nil
			err = nil
		}
		// If we read nothing from the buffer, continue to trying to receive.
		// This saves us when the last Read call emptied the buffer and this
		// call triggers the EOF. There's probably a better way of doing this,
		// but I'm really tired.
		if n > 0 {
			return
		}
	}

	frame, err := c.receiveFrame(true)
	if err != nil {
		return 0, err
	}

	receivedBytes := getFrameData(frame.Data)

	n = copy(b, receivedBytes)
	// Fast path: The entire frame's contents have been copied into b.
	if n >= len(receivedBytes) {
		return
	}

	c.readBuf = bytes.NewReader(receivedBytes[n:])
	return
}

// Write writes the contents of b to the WebSocket using a binary opcode.
func (c *Conn) Write(b []byte) (n int, err error) {
	// []byte is converted to an (U)Int8Array by GopherJS, which fullfils the
	// ArrayBufferView definition.
	err = c.Send(b)
	if err != nil {
		return
	}
	n = len(b)
	return
}

// WriteString writes the contents of s to the WebSocket using a text frame
// opcode.
func (c *Conn) WriteString(s string) (n int, err error) {
	err = c.Send(s)
	if err != nil {
		return
	}
	n = len(s)
	return
}

// BUG(nightexcessive): Conn doesn't currently fulfill the net.Conn interface
// because we can't return net.Addr from Conn.LocalAddr and Conn.RemoteAddr
// because net.init() causes a panic due to attempts to make syscalls.
//
// See: https://github.com/gopherjs/gopherjs/issues/123

// LocalAddr would typically return the local network address, but due to
// limitations in the JavaScript API, it is unable to. Calling this method will
// cause a panic.
func (c *Conn) LocalAddr() net.Addr {
	// BUG(nightexcessive): Conn.LocalAddr() panics because the underlying
	// JavaScript API has no way of figuring out the local address.

	// TODO(nightexcessive): Find a more graceful way to handle this
	panic("we are unable to implement websocket.Conn.LocalAddr() due to limitations in the underlying JavaScript API")
}

// RemoteAddr returns the remote network address, based on
// websocket.WebSocket.URL.
func (c *Conn) RemoteAddr() net.Addr {
	wsURL, err := url.Parse(c.URL)
	if err != nil {
		// TODO(nightexcessive): Should we be panicking for this?
		panic(err)
	}
	return &Addr{wsURL}
}

// SetDeadline sets the read and write deadlines associated with the connection.
// It is equivalent to calling both SetReadDeadline and SetWriteDeadline.
//
// A zero value for t means that I/O operations will not time out.
func (c *Conn) SetDeadline(t time.Time) error {
	c.readDeadline = t
	return nil
}

// SetReadDeadline sets the deadline for future Read calls. A zero value for t
// means Read will not time out.
func (c *Conn) SetReadDeadline(t time.Time) error {
	c.readDeadline = t
	return nil
}

// SetWriteDeadline sets the deadline for future Write calls. Because our writes
// do not block, this function is a no-op.
func (c *Conn) SetWriteDeadline(t time.Time) error {
	return nil
}
<|endoftext|>


<commit_before>package neptulon

import (
	"io"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"

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

// Conn is a client connection.
type Conn struct {
	ID         string
	Session    *cmap.CMap
	middleware []func(ctx *ReqCtx) error
	resRoutes  *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws         *websocket.Conn
	deadline   time.Duration
	closed     bool
}

// NewConn creates a new Conn object.
func NewConn() (*Conn, error) {
	id, err := shortid.UUID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:        id,
		Session:   cmap.New(),
		resRoutes: cmap.New(),
		deadline:  time.Second * time.Duration(300),
	}, nil
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
// Default value for read/write deadline is 300 seconds.
func (c *Conn) SetDeadline(seconds int) {
	c.deadline = time.Second * time.Duration(seconds)
}

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...func(ctx *ReqCtx) error) {
	c.middleware = append(c.middleware, middleware...)
}

// Connect connects to the given WebSocket server.
func (c *Conn) Connect(addr string) error {
	ws, err := websocket.Dial(addr, "", "http://localhost")
	if err != nil {
		return err
	}

	c.ws = ws
	go c.startReceive()
	time.Sleep(time.Millisecond) // give receive goroutine a few cycles to start
	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.ws.RemoteAddr()
}

// SendRequest sends a JSON-RPC request through the connection with an auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequest(method string, params interface{}, resHandler func(res *ResCtx) error) (reqID string, err error) {
	id, err := shortid.UUID()
	if err != nil {
		return "", err
	}

	req := request{ID: id, Method: method, Params: params}
	if err = c.send(req); err != nil {
		return "", err
	}

	c.resRoutes.Set(req.ID, resHandler)
	return id, nil
}

// SendRequestArr sends a JSON-RPC request through the connection, with array params and auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequestArr(method string, resHandler func(res *ResCtx) error, params ...interface{}) (reqID string, err error) {
	return c.SendRequest(method, params, resHandler)
}

// Close closes the connection.
func (c *Conn) Close() error {
	c.closed = true
	return c.ws.Close()
}

// SendResponse sends a JSON-RPC response message through the connection.
func (c *Conn) sendResponse(id string, result interface{}, err *ResError) error {
	return c.send(response{ID: id, Result: result, Error: err})
}

// Send sends the given message through the connection.
func (c *Conn) send(msg interface{}) error {
	if err := c.ws.SetWriteDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Send(c.ws, msg)
}

// Receive receives message from the connection.
func (c *Conn) receive(msg *message) error {
	if err := c.ws.SetReadDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Receive(c.ws, &msg)
}

// UseConn reuses an established websocket.Conn.
// This function blocks and does not return until the connection is closed by another goroutine.
func (c *Conn) useConn(ws *websocket.Conn) {
	c.ws = ws
	c.startReceive()
}

// startReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) startReceive() {
	// append the last middleware to request stack, which will write the response to connection, if any
	c.middleware = append(c.middleware, func(ctx *ReqCtx) error {
		if ctx.Res != nil || ctx.Err != nil {
			return ctx.Conn.sendResponse(ctx.ID, ctx.Res, ctx.Err)
		}

		return nil
	})

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			// if peer closed the connection
			if err == io.EOF {
				log.Printf("Peer disconnected. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			// if we closed the connection
			if c.closed {
				log.Printf("Connection closed. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			log.Println("Error while receiving message:", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
				log.Println("Error while handling request:", err)
				break
			}

			continue
		}

		// if the message is not a JSON-RPC message
		if m.ID == "" || (m.Result == nil && m.Error == nil) {
			log.Printf("Received an unknown message. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
			c.resRoutes.Delete(m.ID)
			if err != nil {
				log.Println("Error while handling response:", err)
				break
			}
		} else {
			log.Println("Error while handling response: got response to a request with unknown ID:", m.ID)
			break
		}
	}

	c.closed = true
}
<commit_msg>Conn.closed -> Conn.connected<commit_after>package neptulon

import (
	"errors"
	"io"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"

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

// Conn is a client connection.
type Conn struct {
	ID           string
	Session      *cmap.CMap
	middleware   []func(ctx *ReqCtx) error
	resRoutes    *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws           *websocket.Conn
	deadline     time.Duration
	isClientConn bool
	running      bool
}

// NewConn creates a new Conn object.
func NewConn() (*Conn, error) {
	id, err := shortid.UUID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:        id,
		Session:   cmap.New(),
		resRoutes: cmap.New(),
		deadline:  time.Second * time.Duration(300),
	}, nil
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
// Default value for read/write deadline is 300 seconds.
func (c *Conn) SetDeadline(seconds int) {
	c.deadline = time.Second * time.Duration(seconds)
}

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...func(ctx *ReqCtx) error) {
	c.middleware = append(c.middleware, middleware...)
}

// Connect connects to the given WebSocket server.
func (c *Conn) Connect(addr string) error {
	ws, err := websocket.Dial(addr, "", "http://localhost")
	if err != nil {
		return err
	}

	c.ws = ws
	c.running = true
	go c.startReceive()
	time.Sleep(time.Millisecond) // give receive goroutine a few cycles to start
	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	if c.ws == nil {
		return nil
	}

	return c.ws.RemoteAddr()
}

// SendRequest sends a JSON-RPC request through the connection with an auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequest(method string, params interface{}, resHandler func(res *ResCtx) error) (reqID string, err error) {
	id, err := shortid.UUID()
	if err != nil {
		return "", err
	}

	req := request{ID: id, Method: method, Params: params}
	if err = c.send(req); err != nil {
		return "", err
	}

	c.resRoutes.Set(req.ID, resHandler)
	return id, nil
}

// SendRequestArr sends a JSON-RPC request through the connection, with array params and auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequestArr(method string, resHandler func(res *ResCtx) error, params ...interface{}) (reqID string, err error) {
	return c.SendRequest(method, params, resHandler)
}

// Close closes the connection.
func (c *Conn) Close() error {
	c.running = false
	return c.ws.Close()
}

// SendResponse sends a JSON-RPC response message through the connection.
func (c *Conn) sendResponse(id string, result interface{}, err *ResError) error {
	return c.send(response{ID: id, Result: result, Error: err})
}

// Send sends the given message through the connection.
func (c *Conn) send(msg interface{}) error {
	if !c.running {
		return errors.New("use of closed connection")
	}

	if err := c.ws.SetWriteDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Send(c.ws, msg)
}

// Receive receives message from the connection.
func (c *Conn) receive(msg *message) error {
	if !c.running {
		return errors.New("use of closed connection")
	}

	if err := c.ws.SetReadDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Receive(c.ws, &msg)
}

// UseConn reuses an established websocket.Conn.
// This function blocks and does not return until the connection is closed by another goroutine.
func (c *Conn) useConn(ws *websocket.Conn) {
	c.ws = ws
	c.running = true
	c.startReceive()
}

// startReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) startReceive() {
	// defer func() {
	// 	if err := recover(); err != nil {
	// 		const size = 64 << 10
	// 		buf := make([]byte, size)
	// 		buf = buf[:runtime.Stack(buf, false)]
	// 		log.Printf("conn: panic serving %v: %v\n%s", c.remoteAddr, err, buf)
	// 	}
	// }()

	// append the last middleware to request stack, which will write the response to connection, if any
	c.middleware = append(c.middleware, func(ctx *ReqCtx) error {
		if ctx.Res != nil || ctx.Err != nil {
			return ctx.Conn.sendResponse(ctx.ID, ctx.Res, ctx.Err)
		}

		return nil
	})

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			// if peer closed the connection
			if err == io.EOF {
				log.Printf("Peer disconnected. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			// if we closed the connection
			if !c.running {
				log.Printf("Connection closed. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			log.Println("Error while receiving message:", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
				log.Println("Error while handling request:", err)
				break
			}

			continue
		}

		// if the message is not a JSON-RPC message
		if m.ID == "" || (m.Result == nil && m.Error == nil) {
			log.Printf("Received an unknown message. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
			c.resRoutes.Delete(m.ID)
			if err != nil {
				log.Println("Error while handling response:", err)
				break
			}
		} else {
			log.Println("Error while handling response: got response to a request with unknown ID:", m.ID)
			break
		}
	}

	c.Close()
}
<|endoftext|>


<commit_before>package turms

import (
	"github.com/ugorji/go/codec"
	"golang.org/x/net/context"
	"net"
	"time"
)

// A Conn represents a connection between two Peers.
type Conn interface {
	Read(context.Context) (Message, error)
	Send(context.Context, Message) error
	Close() error
}

type conn struct {
	conn net.Conn
	mr   *messageReader
	dec  *codec.Decoder
	enc  *codec.Encoder
}

// NewConn wraps a net.Conn
func NewConn(c net.Conn, h codec.Handle) Conn {
	mr := newMessageReader(c, 1024)
	// Force DecodeOptions.ErrorIfNoArrayExpand = false
	// and DecodeOptions.StructToArray = true
	switch handle := h.(type) {
	case *codec.JsonHandle:
		handle.DecodeOptions.ErrorIfNoArrayExpand = false
		handle.EncodeOptions.StructToArray = true
	case *codec.MsgpackHandle:
		handle.DecodeOptions.ErrorIfNoArrayExpand = false
		handle.EncodeOptions.StructToArray = true
	}
	return &conn{
		conn: c,
		mr:   mr,
		dec:  codec.NewDecoder(mr, h),
		enc:  codec.NewEncoder(c, h),
	}
}

func (c *conn) Close() error {
	return c.conn.Close()
}

func (c *conn) Send(ctx context.Context, msg Message) error {
	res := make(chan error, 1)
	go func() {
		err := c.enc.Encode(msg)
		res <- err
	}()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-res:
		return err
	}
}

func (c *conn) Read(ctx context.Context) (Message, error) {
	type msgAndErr struct {
		msg Message
		err error
	}
	res := make(chan msgAndErr, 1)

	go func() {
		var msgTyp [1]MessageType
		err := c.dec.Decode(&msgTyp)
		c.mr.ResetRead()
		if err != nil {
			res <- msgAndErr{msg: nil, err: err}
			return
		}
		msg := NewMessage(msgTyp[0])
		err = c.dec.Decode(msg)
		c.mr.Reset()
		res <- msgAndErr{msg: msg, err: err}
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case r := <-res:
		return r.msg, r.err
	}
}

func waitForMessage(parentCtx context.Context, c Conn, duration time.Duration) (Message, error) {
	ctx, cancel := context.WithTimeout(parentCtx, duration)
	defer cancel()
	return c.Read(ctx)
}
<commit_msg>add two mutex to protect the decoder and encoder of a connection, avoiding data races in both.<commit_after>package turms

import (
	"github.com/ugorji/go/codec"
	"golang.org/x/net/context"
	"net"
	"sync"
	"time"
)

// A Conn represents a connection between two Peers.
type Conn interface {
	Read(context.Context) (Message, error)
	Send(context.Context, Message) error
	Close() error
}

type conn struct {
	conn  net.Conn
	mr    *messageReader
	dec   *codec.Decoder
	decmu sync.Mutex
	enc   *codec.Encoder
	encmu sync.Mutex
}

// NewConn wraps a net.Conn
func NewConn(c net.Conn, h codec.Handle) Conn {
	mr := newMessageReader(c, 1024)
	// Force DecodeOptions.ErrorIfNoArrayExpand = false
	// and DecodeOptions.StructToArray = true
	switch handle := h.(type) {
	case *codec.JsonHandle:
		handle.DecodeOptions.ErrorIfNoArrayExpand = false
		handle.EncodeOptions.StructToArray = true
	case *codec.MsgpackHandle:
		handle.DecodeOptions.ErrorIfNoArrayExpand = false
		handle.EncodeOptions.StructToArray = true
	}
	return &conn{
		conn: c,
		mr:   mr,
		dec:  codec.NewDecoder(mr, h),
		enc:  codec.NewEncoder(c, h),
	}
}

func (c *conn) Close() error {
	return c.conn.Close()
}

func (c *conn) Send(ctx context.Context, msg Message) error {
	res := make(chan error, 1)
	c.encmu.Lock()
	defer c.encmu.Unlock()
	go func() {
		err := c.enc.Encode(msg)
		c.enc.Reset(c.conn)
		res <- err
	}()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-res:
		return err
	}
}

func (c *conn) Read(ctx context.Context) (Message, error) {
	type msgAndErr struct {
		msg Message
		err error
	}
	res := make(chan msgAndErr, 1)
	c.decmu.Lock()
	defer c.decmu.Unlock()

	go func() {
		var msgTyp [1]MessageType
		err := c.dec.Decode(&msgTyp)
		c.mr.ResetRead()
		if err != nil {
			res <- msgAndErr{msg: nil, err: err}
			return
		}
		msg := NewMessage(msgTyp[0])
		err = c.dec.Decode(msg)
		c.mr.Reset()
		c.dec.Reset(c.mr)
		res <- msgAndErr{msg: msg, err: err}
	}()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case r := <-res:
		return r.msg, r.err
	}
}

func waitForMessage(parentCtx context.Context, c Conn, duration time.Duration) (Message, error) {
	ctx, cancel := context.WithTimeout(parentCtx, duration)
	defer cancel()
	return c.Read(ctx)
}
<|endoftext|>


<commit_before>package epp

import (
	"encoding/binary"
	"encoding/xml"
	"io"
	"net"
	"sync"
	"time"
)

// IgnoreEOF returns err unless err == io.EOF,
// in which case it returns nil.
func IgnoreEOF(err error) error {
	if err == io.EOF {
		return nil
	}
	return err
}

// Conn represents a single connection to an EPP server.
// Reads and writes are serialized, so it is safe for concurrent use.
type Conn struct {
	// Conn is the underlying net.Conn (usually a TLS connection).
	net.Conn

	// Timeout defines the timeout for network operations.
	// It must be set at initialization. Changing it after
	// a connection is already opened will have no effect.
	Timeout time.Duration

	// m protects Greeting and LoginResult.
	m sync.Mutex

	// Greeting holds the last received greeting message from the server,
	// indicating server name, status, data policy and capabilities.
	//
	// Deprecated: This field is written to upon opening a new EPP connection and should not be modified.
	Greeting

	// LoginResult holds the last received login response message's Result
	// from the server, in which some servers might include diagnostics such
	// as connection count limits.
	//
	// Deprecated: this field is written to by the Login method but otherwise is not used by this package.
	LoginResult Result

	// mWrite synchronizes connection writes.
	mWrite sync.Mutex

	responses chan *Response
	readErr   error

	done chan struct{}
}

// NewConn initializes an epp.Conn from a net.Conn and performs the EPP
// handshake. It reads and stores the initial EPP <greeting> message.
// https://tools.ietf.org/html/rfc5730#section-2.4
func NewConn(conn net.Conn) (*Conn, error) {
	return NewTimeoutConn(conn, 0)
}

// NewTimeoutConn initializes an epp.Conn like NewConn, limiting the duration of network
// operations on conn using Set(Read|Write)Deadline.
func NewTimeoutConn(conn net.Conn, timeout time.Duration) (*Conn, error) {
	c := newConn(conn)
	c.Timeout = timeout
	g, err := c.readGreeting()
	if err == nil {
		c.m.Lock()
		c.Greeting = g
		c.m.Unlock()
	}
	return c, err
}

// Close sends an EPP <logout> command and closes the connection c.
func (c *Conn) Close() error {
	select {
	case <-c.done:
		return net.ErrClosed
	default:
	}
	c.Logout()
	close(c.done)
	return c.Conn.Close()
}

// newConn initializes an epp.Conn from a net.Conn.
// Used internally for testing.
func newConn(conn net.Conn) *Conn {
	c := Conn{
		Conn:      conn,
		responses: make(chan *Response),
		done:      make(chan struct{}),
	}
	go c.readLoop()
	return &c
}

// writeRequest writes a single EPP request (x) for writing on c.
// writeRequest can be called from multiple goroutines.
func (c *Conn) writeRequest(x []byte) error {
	c.mWrite.Lock()
	defer c.mWrite.Unlock()
	if c.Timeout > 0 {
		c.Conn.SetWriteDeadline(time.Now().Add(c.Timeout))
	}
	return writeDataUnit(c.Conn, x)
}

// readResponse dequeues and returns a EPP response from c.
// It returns an error if the EPP response contains an error Result.
// readResponse can be called from multiple goroutines.
func (c *Conn) readResponse() (*Response, error) {
	select {
	case res := <-c.responses:
		if res == nil {
			return res, c.readErr
		}
		if res.Result.IsError() {
			return nil, &res.Result
		}
		return res, nil
	case <-c.done:
		return nil, net.ErrClosed
	}
}

func (c *Conn) readLoop() {
	defer close(c.responses)
	timeout := c.Timeout
	r := &io.LimitedReader{R: c.Conn}
	decoder := xml.NewDecoder(r)
	for {
		if timeout > 0 {
			c.Conn.SetReadDeadline(time.Now().Add(timeout))
		}
		n, err := parseDataUnit(c.Conn)
		if err != nil {
			c.readErr = err
			return
		}
		r.N = int64(n)
		res := &Response{}
		err = IgnoreEOF(scanResponse.Scan(decoder, res))
		if err != nil {
			c.readErr = err
			return
		}
		c.responses <- res
	}
}

// writeDataUnit writes x to w.
// Bytes written are prefixed with 32-bit header specifying the total size
// of the data unit (message + 4 byte header), in network (big-endian) order.
// http://www.ietf.org/rfc/rfc4934.txt
func writeDataUnit(w io.Writer, x []byte) error {
	logXML("<-- WRITE DATA UNIT -->", x)
	s := uint32(4 + len(x))
	err := binary.Write(w, binary.BigEndian, s)
	if err != nil {
		return err
	}
	_, err = w.Write(x)
	return err
}

// parseDataUnit reads a single EPP data unit header from r, returning the payload size or an error.
// An EPP data unit is prefixed with 32-bit header specifying the total size
// of the data unit (message + 4 byte header), in network (big-endian) order.
// http://www.ietf.org/rfc/rfc4934.txt
func parseDataUnit(r io.Reader) (int32, error) {
	var n int32
	err := binary.Read(r, binary.BigEndian, &n)
	if err != nil {
		return 0, err
	}
	n -= 4 // https://tools.ietf.org/html/rfc5734#section-4
	if n < 0 {
		return 0, io.ErrUnexpectedEOF
	}
	return n, err
}
<commit_msg>int32->uint32 (PR feedback)<commit_after>package epp

import (
	"encoding/binary"
	"encoding/xml"
	"io"
	"net"
	"sync"
	"time"
)

// IgnoreEOF returns err unless err == io.EOF,
// in which case it returns nil.
func IgnoreEOF(err error) error {
	if err == io.EOF {
		return nil
	}
	return err
}

// Conn represents a single connection to an EPP server.
// Reads and writes are serialized, so it is safe for concurrent use.
type Conn struct {
	// Conn is the underlying net.Conn (usually a TLS connection).
	net.Conn

	// Timeout defines the timeout for network operations.
	// It must be set at initialization. Changing it after
	// a connection is already opened will have no effect.
	Timeout time.Duration

	// m protects Greeting and LoginResult.
	m sync.Mutex

	// Greeting holds the last received greeting message from the server,
	// indicating server name, status, data policy and capabilities.
	//
	// Deprecated: This field is written to upon opening a new EPP connection and should not be modified.
	Greeting

	// LoginResult holds the last received login response message's Result
	// from the server, in which some servers might include diagnostics such
	// as connection count limits.
	//
	// Deprecated: this field is written to by the Login method but otherwise is not used by this package.
	LoginResult Result

	// mWrite synchronizes connection writes.
	mWrite sync.Mutex

	responses chan *Response
	readErr   error

	done chan struct{}
}

// NewConn initializes an epp.Conn from a net.Conn and performs the EPP
// handshake. It reads and stores the initial EPP <greeting> message.
// https://tools.ietf.org/html/rfc5730#section-2.4
func NewConn(conn net.Conn) (*Conn, error) {
	return NewTimeoutConn(conn, 0)
}

// NewTimeoutConn initializes an epp.Conn like NewConn, limiting the duration of network
// operations on conn using Set(Read|Write)Deadline.
func NewTimeoutConn(conn net.Conn, timeout time.Duration) (*Conn, error) {
	c := newConn(conn)
	c.Timeout = timeout
	g, err := c.readGreeting()
	if err == nil {
		c.m.Lock()
		c.Greeting = g
		c.m.Unlock()
	}
	return c, err
}

// Close sends an EPP <logout> command and closes the connection c.
func (c *Conn) Close() error {
	select {
	case <-c.done:
		return net.ErrClosed
	default:
	}
	c.Logout()
	close(c.done)
	return c.Conn.Close()
}

// newConn initializes an epp.Conn from a net.Conn.
// Used internally for testing.
func newConn(conn net.Conn) *Conn {
	c := Conn{
		Conn:      conn,
		responses: make(chan *Response),
		done:      make(chan struct{}),
	}
	go c.readLoop()
	return &c
}

// writeRequest writes a single EPP request (x) for writing on c.
// writeRequest can be called from multiple goroutines.
func (c *Conn) writeRequest(x []byte) error {
	c.mWrite.Lock()
	defer c.mWrite.Unlock()
	if c.Timeout > 0 {
		c.Conn.SetWriteDeadline(time.Now().Add(c.Timeout))
	}
	return writeDataUnit(c.Conn, x)
}

// readResponse dequeues and returns a EPP response from c.
// It returns an error if the EPP response contains an error Result.
// readResponse can be called from multiple goroutines.
func (c *Conn) readResponse() (*Response, error) {
	select {
	case res := <-c.responses:
		if res == nil {
			return res, c.readErr
		}
		if res.Result.IsError() {
			return nil, &res.Result
		}
		return res, nil
	case <-c.done:
		return nil, net.ErrClosed
	}
}

func (c *Conn) readLoop() {
	defer close(c.responses)
	timeout := c.Timeout
	r := &io.LimitedReader{R: c.Conn}
	decoder := xml.NewDecoder(r)
	for {
		if timeout > 0 {
			c.Conn.SetReadDeadline(time.Now().Add(timeout))
		}
		n, err := parseDataUnit(c.Conn)
		if err != nil {
			c.readErr = err
			return
		}
		r.N = int64(n)
		res := &Response{}
		err = IgnoreEOF(scanResponse.Scan(decoder, res))
		if err != nil {
			c.readErr = err
			return
		}
		c.responses <- res
	}
}

// writeDataUnit writes x to w.
// Bytes written are prefixed with 32-bit header specifying the total size
// of the data unit (message + 4 byte header), in network (big-endian) order.
// http://www.ietf.org/rfc/rfc4934.txt
func writeDataUnit(w io.Writer, x []byte) error {
	logXML("<-- WRITE DATA UNIT -->", x)
	s := uint32(4 + len(x))
	err := binary.Write(w, binary.BigEndian, s)
	if err != nil {
		return err
	}
	_, err = w.Write(x)
	return err
}

// parseDataUnit reads a single EPP data unit header from r, returning the payload size or an error.
// An EPP data unit is prefixed with 32-bit header specifying the total size
// of the data unit (message + 4 byte header), in network (big-endian) order.
// http://www.ietf.org/rfc/rfc4934.txt
func parseDataUnit(r io.Reader) (uint32, error) {
	var n uint32
	err := binary.Read(r, binary.BigEndian, &n)
	if err != nil {
		return 0, err
	}
	n -= 4 // https://tools.ietf.org/html/rfc5734#section-4
	if n < 0 {
		return 0, io.ErrUnexpectedEOF
	}
	return n, err
}
<|endoftext|>


<commit_before>package ucon

import (
	"fmt"
	"net/http"
	"strings"

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

// DefaultMux is the default ServeMux in ucon.
var DefaultMux = NewServeMux()

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux {
	mux := &ServeMux{
		router: &Router{},
	}
	mux.router.mux = mux
	return mux
}

// ServeMux is an HTTP request multiplexer.
type ServeMux struct {
	Debug       bool
	router      *Router
	middlewares []MiddlewareFunc
	plugins     []*pluginContainer
}

// MiddlewareFunc is an adapter to hook middleware processing.
// Middlrware works with 1 request.
type MiddlewareFunc func(b *Bubble) error

// Middleware can append Middleware to ServeMux.
func (m *ServeMux) Middleware(f MiddlewareFunc) {
	m.middlewares = append(m.middlewares, f)
}

// Plugin can append Plugin to ServeMux.
func (m *ServeMux) Plugin(plugin interface{}) {
	p, ok := plugin.(*pluginContainer)
	if !ok {
		p = &pluginContainer{base: plugin}
	}
	p.check()
	m.plugins = append(m.plugins, p)
}

// Prepare the ServeMux.
// Plugin is not show affect to anything.
// This method is enabled plugins.
func (m *ServeMux) Prepare() {
	for _, plugin := range m.plugins {
		used := false
		if sc := plugin.HandlersScanner(); sc != nil {
			err := sc.HandlersScannerProcess(m, m.router.handlers)
			if err != nil {
				panic(err)
			}
			used = true
		}
		if !used {
			panic(fmt.Sprintf("unused plugin: %#v", plugin))
		}
	}
}

// ServeHTTP dispatches request to the handler.
func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// NOTE ucon内部のルーティングは一元的にこの関数から行う
	// Handlerを細分化しhttp.ServeMuxに登録すると、OPTIONSのhandleがうまくできなくなる
	// このため、Handlerはucon全体で1つとし、OPTIONSも通常のMethodと同じようにHandlerを設定し利用する
	// OPTIONSを適切にhandleするため、全てのHandlerに特殊なHookを入れるよりマシである

	m.router.ServeHTTP(w, r)
}

// ListenAndServe start accepts the client request.
func (m *ServeMux) ListenAndServe(addr string) error {
	m.Prepare()

	server := &http.Server{Addr: addr, Handler: m}
	return server.ListenAndServe()
}

// Handle register the HandlerContainer for the given method & path to the ServeMux.
func (m *ServeMux) Handle(method string, path string, hc HandlerContainer) {
	CheckFunction(hc.Handler())

	pathTmpl := ParsePathTemplate(path)
	methods := strings.Split(strings.ToUpper(method), ",")
	for _, method := range methods {
		rd := &RouteDefinition{
			Method:           method,
			PathTemplate:     pathTmpl,
			HandlerContainer: hc,
		}
		m.router.addRoute(rd)
	}
}

// HandleFunc register the handler function for the given method & path to the ServeMux.
func (m *ServeMux) HandleFunc(method string, path string, h interface{}) {
	m.Handle(method, path, &handlerContainerImpl{
		handler: h,
		Context: background,
	})
}

func (m *ServeMux) newBubble(c context.Context, w http.ResponseWriter, r *http.Request, rd *RouteDefinition) (*Bubble, error) {
	b := &Bubble{
		R:              r,
		W:              w,
		Context:        c,
		RequestHandler: rd.HandlerContainer,
	}
	err := b.init(m)
	if err != nil {
		return nil, err
	}

	return b, nil
}

// HandlerContainer is handler function container.
// and It has a ucon Context that make it possible communicate to Plugins.
type HandlerContainer interface {
	Handler() interface{}
	Context
}

type handlerContainerImpl struct {
	handler interface{}
	Context
}

func (hc *handlerContainerImpl) Handler() interface{} {
	return hc.handler
}

// Orthodox middlewares enable to DefaultServeMux.
func Orthodox() {
	DefaultMux.Middleware(ResponseMapper())
	DefaultMux.Middleware(HTTPRWDI())
	DefaultMux.Middleware(NetContextDI())
	DefaultMux.Middleware(RequestObjectMapper())
}

// Middleware can append Middleware to ServeMux.
func Middleware(f MiddlewareFunc) {
	DefaultMux.Middleware(f)
}

// Plugin can append Plugin to ServeMux.
func Plugin(plugin interface{}) {
	DefaultMux.Plugin(plugin)
}

// ListenAndServe start accepts the client request.
func ListenAndServe(addr string) {
	DefaultMux.ListenAndServe(addr)
}

// Handle register the HandlerContainer for the given method & path to the ServeMux.
func Handle(method string, path string, hc HandlerContainer) {
	DefaultMux.Handle(method, path, hc)
}

// HandleFunc register the handler function for the given method & path to the ServeMux.
func HandleFunc(method string, path string, h interface{}) {
	DefaultMux.HandleFunc(method, path, h)
}
<commit_msg>fix typo<commit_after>package ucon

import (
	"fmt"
	"net/http"
	"strings"

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

// DefaultMux is the default ServeMux in ucon.
var DefaultMux = NewServeMux()

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux {
	mux := &ServeMux{
		router: &Router{},
	}
	mux.router.mux = mux
	return mux
}

// ServeMux is an HTTP request multiplexer.
type ServeMux struct {
	Debug       bool
	router      *Router
	middlewares []MiddlewareFunc
	plugins     []*pluginContainer
}

// MiddlewareFunc is an adapter to hook middleware processing.
// Middleware works with 1 request.
type MiddlewareFunc func(b *Bubble) error

// Middleware can append Middleware to ServeMux.
func (m *ServeMux) Middleware(f MiddlewareFunc) {
	m.middlewares = append(m.middlewares, f)
}

// Plugin can append Plugin to ServeMux.
func (m *ServeMux) Plugin(plugin interface{}) {
	p, ok := plugin.(*pluginContainer)
	if !ok {
		p = &pluginContainer{base: plugin}
	}
	p.check()
	m.plugins = append(m.plugins, p)
}

// Prepare the ServeMux.
// Plugin is not show affect to anything.
// This method is enabled plugins.
func (m *ServeMux) Prepare() {
	for _, plugin := range m.plugins {
		used := false
		if sc := plugin.HandlersScanner(); sc != nil {
			err := sc.HandlersScannerProcess(m, m.router.handlers)
			if err != nil {
				panic(err)
			}
			used = true
		}
		if !used {
			panic(fmt.Sprintf("unused plugin: %#v", plugin))
		}
	}
}

// ServeHTTP dispatches request to the handler.
func (m *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// NOTE ucon内部のルーティングは一元的にこの関数から行う
	// Handlerを細分化しhttp.ServeMuxに登録すると、OPTIONSのhandleがうまくできなくなる
	// このため、Handlerはucon全体で1つとし、OPTIONSも通常のMethodと同じようにHandlerを設定し利用する
	// OPTIONSを適切にhandleするため、全てのHandlerに特殊なHookを入れるよりマシである

	m.router.ServeHTTP(w, r)
}

// ListenAndServe start accepts the client request.
func (m *ServeMux) ListenAndServe(addr string) error {
	m.Prepare()

	server := &http.Server{Addr: addr, Handler: m}
	return server.ListenAndServe()
}

// Handle register the HandlerContainer for the given method & path to the ServeMux.
func (m *ServeMux) Handle(method string, path string, hc HandlerContainer) {
	CheckFunction(hc.Handler())

	pathTmpl := ParsePathTemplate(path)
	methods := strings.Split(strings.ToUpper(method), ",")
	for _, method := range methods {
		rd := &RouteDefinition{
			Method:           method,
			PathTemplate:     pathTmpl,
			HandlerContainer: hc,
		}
		m.router.addRoute(rd)
	}
}

// HandleFunc register the handler function for the given method & path to the ServeMux.
func (m *ServeMux) HandleFunc(method string, path string, h interface{}) {
	m.Handle(method, path, &handlerContainerImpl{
		handler: h,
		Context: background,
	})
}

func (m *ServeMux) newBubble(c context.Context, w http.ResponseWriter, r *http.Request, rd *RouteDefinition) (*Bubble, error) {
	b := &Bubble{
		R:              r,
		W:              w,
		Context:        c,
		RequestHandler: rd.HandlerContainer,
	}
	err := b.init(m)
	if err != nil {
		return nil, err
	}

	return b, nil
}

// HandlerContainer is handler function container.
// and It has a ucon Context that make it possible communicate to Plugins.
type HandlerContainer interface {
	Handler() interface{}
	Context
}

type handlerContainerImpl struct {
	handler interface{}
	Context
}

func (hc *handlerContainerImpl) Handler() interface{} {
	return hc.handler
}

// Orthodox middlewares enable to DefaultServeMux.
func Orthodox() {
	DefaultMux.Middleware(ResponseMapper())
	DefaultMux.Middleware(HTTPRWDI())
	DefaultMux.Middleware(NetContextDI())
	DefaultMux.Middleware(RequestObjectMapper())
}

// Middleware can append Middleware to ServeMux.
func Middleware(f MiddlewareFunc) {
	DefaultMux.Middleware(f)
}

// Plugin can append Plugin to ServeMux.
func Plugin(plugin interface{}) {
	DefaultMux.Plugin(plugin)
}

// ListenAndServe start accepts the client request.
func ListenAndServe(addr string) {
	DefaultMux.ListenAndServe(addr)
}

// Handle register the HandlerContainer for the given method & path to the ServeMux.
func Handle(method string, path string, hc HandlerContainer) {
	DefaultMux.Handle(method, path, hc)
}

// HandleFunc register the handler function for the given method & path to the ServeMux.
func HandleFunc(method string, path string, h interface{}) {
	DefaultMux.HandleFunc(method, path, h)
}
<|endoftext|>


<commit_before>package mango

import (
	"net/http"
	"strconv"
	"strings"
)

// CORSConfig holds CORS configuration. It can be used as
// the configuration for an individual resource or as a
// global configuration for the entire router tree.
type CORSConfig struct {
	// Origins list all permitted origins. A CORS request
	// origin MUST be in the Origins list for the response
	// headers to be populated with the correct response.
	// Values must contain the scheme (e.g. http://here.com).
	// A wildcard * can be used which will match ALL origins,
	// however the Access-Control-Allow-Origin response header
	// will always echo the request Origin if the remaining CORS
	// criteria is met.
	Origins []string
	// Methods available for the resource. If there are methods
	// listed here for which there is no handler, then that
	// method will not be included in the Access-Control-Allow-Methods
	// response header.
	Methods []string
	// Headers lists the custom headers in a request that the
	// server will accept
	Headers []string
	// ExposedHeaders are custom headers that the client browser
	// is allowed to access
	ExposedHeaders []string
	// Credentials controls the Access-Control-Allow-Credentials header.
	// The header is only included in the response if Credentials is true,
	// in which case the header has a value of "true".
	// A value of true allows the client browser to access response cookies.
	Credentials bool
	// MaxAge is the cache duration (in seconds) that is returned
	// in a Preflight Access-Control-Max-Age response header.
	// A value of zero means the header won't be sent.
	MaxAge int
}

func (c *CORSConfig) originAllowed(origin string) bool {
	return stringInSlice("*", c.Origins) ||
		stringInSlice(origin, c.Origins)
}

func (c *CORSConfig) methodAllowed(method string) bool {
	return stringInSlice(method, c.Methods)
}

func (c *CORSConfig) headersAllowed(headers string) bool {
	ha := strings.Split(headers, ",")
	for i := 0; i < len(ha); i++ {
		h := strings.TrimSpace(ha[i])
		if !stringInSlice(h, c.Headers) {
			return false
		}
	}
	return true
}

func (c *CORSConfig) clone() *CORSConfig {
	cl := CORSConfig{
		Origins:        c.Origins,
		Methods:        c.Methods,
		Headers:        c.Headers,
		ExposedHeaders: c.ExposedHeaders,
		Credentials:    c.Credentials,
		MaxAge:         c.MaxAge,
	}
	return &cl
}

func (c *CORSConfig) merge(m CORSConfig) {
	c.Origins = appendIfNotExists(c.Origins, m.Origins)
	c.Methods = appendIfNotExists(c.Methods, m.Methods)
	c.Headers = appendIfNotExists(c.Headers, m.Headers)
	c.ExposedHeaders = appendIfNotExists(c.ExposedHeaders, m.ExposedHeaders)
	c.Credentials = m.Credentials
	c.MaxAge = m.MaxAge
}

func appendIfNotExists(dest []string, src []string) []string {
	for _, v := range src {
		if !stringInSlice(v, dest) {
			dest = append(dest, v)
		}
	}
	return dest
}

func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

// type CORSType int
//
// const (
// 	NoCORS CORSType = 0 + iota
// 	SimpleCORS
// 	PreflightCORS
// )

func handleCORS(req *http.Request, w http.ResponseWriter, resource *Resource) (preflight bool) {
	origin := req.Header.Get("Origin")
	corsConf := (*resource).CORSConfig
	if corsConf == nil {
		return
	}
	if !(*corsConf).originAllowed(origin) {
		return
	}

	if req.Method == "OPTIONS" {
		// check for preflight
		method := req.Header.Get("Access-Control-Request-Method")
		if method == "" {
			return
		}
		if !corsConf.methodAllowed(method) {
			return
		}
		if _, ok := resource.Handlers[method]; !ok {
			return
		}
		reqHeaders := req.Header.Get("Access-Control-Request-Headers")
		if reqHeaders == "" {
			return
		}
		if !corsConf.headersAllowed(reqHeaders) {
			return
		}

		for _, m := range (*corsConf).Methods {
			if _, ok := resource.Handlers[m]; !ok {
				continue
			}
			w.Header().Add("Access-Control-Allow-Methods", m)
		}

		for _, h := range (*corsConf).Headers {
			w.Header().Add("Access-Control-Allow-Headers", h)
		}
		if (*corsConf).MaxAge > 0 {
			maStr := strconv.Itoa((*corsConf).MaxAge)
			w.Header().Set("Access-Control-Max-Age", maStr)
		}
		w.WriteHeader(http.StatusOK)
		preflight = true
	} else {
		// normal request
		for _, h := range (*corsConf).ExposedHeaders {
			w.Header().Add("Access-Control-Expose-Headers", h)
		}
	}
	w.Header().Set("Access-Control-Allow-Origin", origin)
	w.Header().Add("Vary", "Origin")
	if (*corsConf).Credentials {
		w.Header().Set("Access-Control-Allow-Credentials", "true")
	}
	return
}
<commit_msg>Fix CORS pre-flight request handler<commit_after>package mango

import (
	"net/http"
	"strconv"
	"strings"
)

// CORSConfig holds CORS configuration. It can be used as
// the configuration for an individual resource or as a
// global configuration for the entire router tree.
type CORSConfig struct {
	// Origins list all permitted origins. A CORS request
	// origin MUST be in the Origins list for the response
	// headers to be populated with the correct response.
	// Values must contain the scheme (e.g. http://here.com).
	// A wildcard * can be used which will match ALL origins,
	// however the Access-Control-Allow-Origin response header
	// will always echo the request Origin if the remaining CORS
	// criteria is met.
	Origins []string
	// Methods available for the resource. If there are methods
	// listed here for which there is no handler, then that
	// method will not be included in the Access-Control-Allow-Methods
	// response header.
	Methods []string
	// Headers lists the custom headers in a request that the
	// server will accept
	Headers []string
	// ExposedHeaders are custom headers that the client browser
	// is allowed to access
	ExposedHeaders []string
	// Credentials controls the Access-Control-Allow-Credentials header.
	// The header is only included in the response if Credentials is true,
	// in which case the header has a value of "true".
	// A value of true allows the client browser to access response cookies.
	Credentials bool
	// MaxAge is the cache duration (in seconds) that is returned
	// in a Preflight Access-Control-Max-Age response header.
	// A value of zero means the header won't be sent.
	MaxAge int
}

func (c *CORSConfig) originAllowed(origin string) bool {
	return stringInSlice("*", c.Origins) ||
		stringInSlice(origin, c.Origins)
}

func (c *CORSConfig) methodAllowed(method string) bool {
	return stringInSlice(method, c.Methods)
}

func (c *CORSConfig) headersAllowed(headers string) bool {
	ha := strings.Split(headers, ",")
	for i := 0; i < len(ha); i++ {
		h := strings.TrimSpace(ha[i])
		if !stringInSlice(h, c.Headers) {
			return false
		}
	}
	return true
}

func (c *CORSConfig) clone() *CORSConfig {
	cl := CORSConfig{
		Origins:        c.Origins,
		Methods:        c.Methods,
		Headers:        c.Headers,
		ExposedHeaders: c.ExposedHeaders,
		Credentials:    c.Credentials,
		MaxAge:         c.MaxAge,
	}
	return &cl
}

func (c *CORSConfig) merge(m CORSConfig) {
	c.Origins = appendIfNotExists(c.Origins, m.Origins)
	c.Methods = appendIfNotExists(c.Methods, m.Methods)
	c.Headers = appendIfNotExists(c.Headers, m.Headers)
	c.ExposedHeaders = appendIfNotExists(c.ExposedHeaders, m.ExposedHeaders)
	c.Credentials = m.Credentials
	c.MaxAge = m.MaxAge
}

func appendIfNotExists(dest []string, src []string) []string {
	for _, v := range src {
		if !stringInSlice(v, dest) {
			dest = append(dest, v)
		}
	}
	return dest
}

func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

// type CORSType int
//
// const (
// 	NoCORS CORSType = 0 + iota
// 	SimpleCORS
// 	PreflightCORS
// )

func handleCORS(req *http.Request, w http.ResponseWriter, resource *Resource) (preflight bool) {
	origin := req.Header.Get("Origin")
	corsConf := (*resource).CORSConfig
	if corsConf == nil {
		return
	}
	if !(*corsConf).originAllowed(origin) {
		return
	}

	if req.Method == "OPTIONS" {
		// check for preflight
		method := req.Header.Get("Access-Control-Request-Method")
		if method == "" {
			return
		}
		if !corsConf.methodAllowed(method) {
			return
		}
		if _, ok := resource.Handlers[method]; !ok {
			return
		}
		reqHeaders := req.Header.Get("Access-Control-Request-Headers")
		if reqHeaders == "" {
			return
		}
		if !corsConf.headersAllowed(reqHeaders) {
			return
		}

		for _, m := range (*corsConf).Methods {
			if _, ok := resource.Handlers[m]; !ok {
				continue
			}
			w.Header().Add("Access-Control-Allow-Methods", m)
		}

		for _, h := range (*corsConf).Headers {
			w.Header().Add("Access-Control-Allow-Headers", h)
		}
		if (*corsConf).MaxAge > 0 {
			maStr := strconv.Itoa((*corsConf).MaxAge)
			w.Header().Set("Access-Control-Max-Age", maStr)
		}

		preflight = true
	} else {
		// normal request
		for _, h := range (*corsConf).ExposedHeaders {
			w.Header().Add("Access-Control-Expose-Headers", h)
		}
	}
	w.Header().Set("Access-Control-Allow-Origin", origin)
	w.Header().Add("Vary", "Origin")
	if (*corsConf).Credentials {
		w.Header().Set("Access-Control-Allow-Credentials", "true")
	}
	if(preflight){
		w.WriteHeader(http.StatusOK)
	}
	return
}
<|endoftext|>