<commit_before>// jwe implements JWE https://tools.ietf.org/html/rfc7516

package jwe

import (
	"bytes"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"log"

	"github.com/lestrrat/go-jwx/buffer"
	"github.com/lestrrat/go-jwx/jwa"
)

func debug(f string, args ...interface{}) {
	log.Printf(f, args...)
}

func Parse(buf []byte) (*Message, error) {
	buf = bytes.TrimSpace(buf)
	if len(buf) == 0 {
		return nil, errors.New("empty buffer")
	}

	if buf[0] == '{' {
		return parseJSON(buf)
	}
	return parseCompact(buf)
}

func ParseString(s string) (*Message, error) {
	return Parse([]byte(s))
}

func parseJSON(buf []byte) (*Message, error) {
	m := struct {
		*Message
		*Recipient
	}{}

	if err := json.Unmarshal(buf, &m); err != nil {
		return nil, err
	}

	// if the "signature" field exist, treat it as a flattened
	if m.Recipient != nil {
		if len(m.Message.Recipients) != 0 {
			return nil, errors.New("invalid message: mixed flattened/full json serialization")
		}

		m.Message.Recipients = []Recipient{*m.Recipient}
	}

	return m.Message, nil
}

func parseCompact(buf []byte) (*Message, error) {
	parts := bytes.Split(buf, []byte{'.'})
	if len(parts) != 5 {
		return nil, ErrInvalidCompactPartsCount
	}

	enc := base64.RawURLEncoding
	p0Len := enc.DecodedLen(len(parts[0]))
	p1Len := enc.DecodedLen(len(parts[1]))
	p2Len := enc.DecodedLen(len(parts[2]))
	p3Len := enc.DecodedLen(len(parts[3]))
	p4Len := enc.DecodedLen(len(parts[4]))

	out := make([]byte, p0Len+p1Len+p2Len+p3Len+p4Len)

	hdrbuf := buffer.Buffer(out[:p0Len])
	if _, err := enc.Decode(hdrbuf, parts[0]); err != nil {
		return nil, err
	}
	hdrbuf = bytes.TrimRight(hdrbuf, "\x00")

	hdr := NewHeader()
	if err := json.Unmarshal(hdrbuf, hdr); err != nil {
		return nil, err
	}

	// We need the protected header to contain the content encryption
	// algorithm. XXX probably other headers need to go there too
	protected := NewEncodedHeader()
	protected.ContentEncryption = hdr.ContentEncryption
	hdr.ContentEncryption = ""

	enckeybuf := buffer.Buffer(out[p0Len : p0Len+p1Len])
	if _, err := enc.Decode(enckeybuf, parts[1]); err != nil {
		return nil, err
	}
	enckeybuf = bytes.TrimRight(enckeybuf, "\x00")

	ivbuf := buffer.Buffer(out[p0Len+p1Len : p0Len+p1Len+p2Len])
	if _, err := enc.Decode(ivbuf, parts[2]); err != nil {
		return nil, err
	}
	ivbuf = bytes.TrimRight(ivbuf, "\x00")

	ctbuf := buffer.Buffer(out[p0Len+p1Len+p2Len : p0Len+p1Len+p2Len+p3Len])
	if _, err := enc.Decode(ctbuf, parts[3]); err != nil {
		return nil, err
	}
	ctbuf = bytes.TrimRight(ctbuf, "\x00")

	tagbuf := buffer.Buffer(out[p0Len+p1Len+p2Len+p3Len : p0Len+p1Len+p2Len+p3Len+p4Len])
	if _, err := enc.Decode(tagbuf, parts[4]); err != nil {
		return nil, err
	}
	tagbuf = bytes.TrimRight(tagbuf, "\x00")

	b64hdrbuf, err := hdrbuf.Base64Encode()
	if err != nil {
		return nil, err
	}

	m := NewMessage()
	m.AuthenticatedData = b64hdrbuf
	m.ProtectedHeader = protected
	m.Tag = tagbuf
	m.CipherText = ctbuf
	m.InitializationVector = ivbuf
	m.Recipients = []Recipient{
		Recipient{
			Header:       hdr,
			EncryptedKey: enckeybuf,
		},
	}
	return m, nil
}

// BuildKeyDecrypter creates a new KeyDecrypter instance from the given
// parameters. It is used by the DecryptMessage fucntion to create
// key decrypter(s) from the given message. `keysize` is only used by
// some decrypters. Use the value from ContentCipher.KeySize().
func BuildKeyDecrypter(alg jwa.KeyEncryptionAlgorithm, key interface{}, keysize int) (KeyDecrypter, error) {
	switch alg {
	case jwa.RSA1_5:
		privkey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("*rsa.PrivateKey is required as the key to build this key decrypter")
		}
		return NewRSAPKCS15KeyDecrypt(alg, privkey, keysize), nil
	case jwa.RSA_OAEP, jwa.RSA_OAEP_256:
		privkey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("*rsa.PrivateKey is required as the key to build this key decrypter")
		}
		return NewRSAOAEPKeyDecrypt(alg, privkey), nil
	}

	return nil, ErrUnsupportedAlgorithm
}

func BuildContentCipher(alg jwa.ContentEncryptionAlgorithm) (ContentCipher, error) {
	switch alg {
	case jwa.A128GCM, jwa.A192GCM, jwa.A256GCM, jwa.A128CBC_HS256, jwa.A192CBC_HS384, jwa.A256CBC_HS512:
		return NewAesContentCipher(alg)
	}

	return nil, ErrUnsupportedAlgorithm
}

func DecryptMessage(m *Message, key interface{}) ([]byte, error) {
	var err error

	if len(m.Recipients) == 0 {
		return nil, errors.New("no recipients, can not proceed with decrypt")
	}

	enc := m.ProtectedHeader.ContentEncryption

	h := NewHeader()
	if err := h.Copy(m.ProtectedHeader.Header); err != nil {
		return nil, err
	}
	h, err = h.Merge(m.UnprotectedHeader)
	if err != nil {
		debug("failed to merge unprotected header")
		return nil, err
	}

	aad := m.AuthenticatedData.Bytes()
	ciphertext := m.CipherText.Bytes()
	iv := m.InitializationVector.Bytes()
	tag := m.Tag.Bytes()

	cipher, err := BuildContentCipher(enc)
	if err != nil {
		return nil, fmt.Errorf("unsupported content cipher algorithm '%s'", enc)
	}
	keysize := cipher.KeySize()

	var plaintext []byte
	for _, recipient := range m.Recipients {
		h2 := NewHeader()
		if err := h2.Copy(h); err != nil {
			debug("failed to copy header: %s", err)
			continue
		}

		h2, err := h2.Merge(recipient.Header)
		if err != nil {
			debug("Failed to merge! %s", err)
			continue
		}

		k, err := BuildKeyDecrypter(h2.Algorithm, key, keysize)
		if err != nil {
			debug("failed to crete key decrypter: %s", err)
			continue
		}

		cek, err := k.KeyDecrypt(recipient.EncryptedKey)
		if err != nil {
			debug("failed to decrypt key: %s", err)
			continue
		}

		debug("cek        = %v", cek)
		debug("iv	       = %v", iv)
		debug("ciphertext = %v", ciphertext)
		debug("tag        = %v", tag)
		debug("aad        = %v", aad)
		plaintext, err = cipher.decrypt(cek, iv, ciphertext, tag, aad)
		if err == nil {
			break
		}
	}

	if plaintext == nil {
		return nil, errors.New("failed to decrypt key")
	}

	if h.Compression != "" {
		panic("compression not implemented")
	}

	return plaintext, nil
}
<commit_msg>implement deflate<commit_after>// jwe implements JWE https://tools.ietf.org/html/rfc7516

package jwe

import (
	"bytes"
	"compress/flate"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"log"

	"github.com/lestrrat/go-jwx/buffer"
	"github.com/lestrrat/go-jwx/jwa"
)

func debug(f string, args ...interface{}) {
	log.Printf(f, args...)
}

func Parse(buf []byte) (*Message, error) {
	buf = bytes.TrimSpace(buf)
	if len(buf) == 0 {
		return nil, errors.New("empty buffer")
	}

	if buf[0] == '{' {
		return parseJSON(buf)
	}
	return parseCompact(buf)
}

func ParseString(s string) (*Message, error) {
	return Parse([]byte(s))
}

func parseJSON(buf []byte) (*Message, error) {
	m := struct {
		*Message
		*Recipient
	}{}

	if err := json.Unmarshal(buf, &m); err != nil {
		return nil, err
	}

	// if the "signature" field exist, treat it as a flattened
	if m.Recipient != nil {
		if len(m.Message.Recipients) != 0 {
			return nil, errors.New("invalid message: mixed flattened/full json serialization")
		}

		m.Message.Recipients = []Recipient{*m.Recipient}
	}

	return m.Message, nil
}

func parseCompact(buf []byte) (*Message, error) {
	parts := bytes.Split(buf, []byte{'.'})
	if len(parts) != 5 {
		return nil, ErrInvalidCompactPartsCount
	}

	enc := base64.RawURLEncoding
	p0Len := enc.DecodedLen(len(parts[0]))
	p1Len := enc.DecodedLen(len(parts[1]))
	p2Len := enc.DecodedLen(len(parts[2]))
	p3Len := enc.DecodedLen(len(parts[3]))
	p4Len := enc.DecodedLen(len(parts[4]))

	out := make([]byte, p0Len+p1Len+p2Len+p3Len+p4Len)

	hdrbuf := buffer.Buffer(out[:p0Len])
	if _, err := enc.Decode(hdrbuf, parts[0]); err != nil {
		return nil, err
	}
	hdrbuf = bytes.TrimRight(hdrbuf, "\x00")

	hdr := NewHeader()
	if err := json.Unmarshal(hdrbuf, hdr); err != nil {
		return nil, err
	}

	// We need the protected header to contain the content encryption
	// algorithm. XXX probably other headers need to go there too
	protected := NewEncodedHeader()
	protected.ContentEncryption = hdr.ContentEncryption
	hdr.ContentEncryption = ""

	enckeybuf := buffer.Buffer(out[p0Len : p0Len+p1Len])
	if _, err := enc.Decode(enckeybuf, parts[1]); err != nil {
		return nil, err
	}
	enckeybuf = bytes.TrimRight(enckeybuf, "\x00")

	ivbuf := buffer.Buffer(out[p0Len+p1Len : p0Len+p1Len+p2Len])
	if _, err := enc.Decode(ivbuf, parts[2]); err != nil {
		return nil, err
	}
	ivbuf = bytes.TrimRight(ivbuf, "\x00")

	ctbuf := buffer.Buffer(out[p0Len+p1Len+p2Len : p0Len+p1Len+p2Len+p3Len])
	if _, err := enc.Decode(ctbuf, parts[3]); err != nil {
		return nil, err
	}
	ctbuf = bytes.TrimRight(ctbuf, "\x00")

	tagbuf := buffer.Buffer(out[p0Len+p1Len+p2Len+p3Len : p0Len+p1Len+p2Len+p3Len+p4Len])
	if _, err := enc.Decode(tagbuf, parts[4]); err != nil {
		return nil, err
	}
	tagbuf = bytes.TrimRight(tagbuf, "\x00")

	b64hdrbuf, err := hdrbuf.Base64Encode()
	if err != nil {
		return nil, err
	}

	m := NewMessage()
	m.AuthenticatedData = b64hdrbuf
	m.ProtectedHeader = protected
	m.Tag = tagbuf
	m.CipherText = ctbuf
	m.InitializationVector = ivbuf
	m.Recipients = []Recipient{
		Recipient{
			Header:       hdr,
			EncryptedKey: enckeybuf,
		},
	}
	return m, nil
}

// BuildKeyDecrypter creates a new KeyDecrypter instance from the given
// parameters. It is used by the DecryptMessage fucntion to create
// key decrypter(s) from the given message. `keysize` is only used by
// some decrypters. Use the value from ContentCipher.KeySize().
func BuildKeyDecrypter(alg jwa.KeyEncryptionAlgorithm, key interface{}, keysize int) (KeyDecrypter, error) {
	switch alg {
	case jwa.RSA1_5:
		privkey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("*rsa.PrivateKey is required as the key to build this key decrypter")
		}
		return NewRSAPKCS15KeyDecrypt(alg, privkey, keysize), nil
	case jwa.RSA_OAEP, jwa.RSA_OAEP_256:
		privkey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("*rsa.PrivateKey is required as the key to build this key decrypter")
		}
		return NewRSAOAEPKeyDecrypt(alg, privkey), nil
	}

	return nil, ErrUnsupportedAlgorithm
}

func BuildContentCipher(alg jwa.ContentEncryptionAlgorithm) (ContentCipher, error) {
	switch alg {
	case jwa.A128GCM, jwa.A192GCM, jwa.A256GCM, jwa.A128CBC_HS256, jwa.A192CBC_HS384, jwa.A256CBC_HS512:
		return NewAesContentCipher(alg)
	}

	return nil, ErrUnsupportedAlgorithm
}

func DecryptMessage(m *Message, key interface{}) ([]byte, error) {
	var err error

	if len(m.Recipients) == 0 {
		return nil, errors.New("no recipients, can not proceed with decrypt")
	}

	enc := m.ProtectedHeader.ContentEncryption

	h := NewHeader()
	if err := h.Copy(m.ProtectedHeader.Header); err != nil {
		return nil, err
	}
	h, err = h.Merge(m.UnprotectedHeader)
	if err != nil {
		debug("failed to merge unprotected header")
		return nil, err
	}

	aad := m.AuthenticatedData.Bytes()
	ciphertext := m.CipherText.Bytes()
	iv := m.InitializationVector.Bytes()
	tag := m.Tag.Bytes()

	cipher, err := BuildContentCipher(enc)
	if err != nil {
		return nil, fmt.Errorf("unsupported content cipher algorithm '%s'", enc)
	}
	keysize := cipher.KeySize()

	var plaintext []byte
	for _, recipient := range m.Recipients {
		h2 := NewHeader()
		if err := h2.Copy(h); err != nil {
			debug("failed to copy header: %s", err)
			continue
		}

		h2, err := h2.Merge(recipient.Header)
		if err != nil {
			debug("Failed to merge! %s", err)
			continue
		}

		k, err := BuildKeyDecrypter(h2.Algorithm, key, keysize)
		if err != nil {
			debug("failed to crete key decrypter: %s", err)
			continue
		}

		cek, err := k.KeyDecrypt(recipient.EncryptedKey)
		if err != nil {
			debug("failed to decrypt key: %s", err)
			continue
		}

		debug("cek        = %v", cek)
		debug("iv	       = %v", iv)
		debug("ciphertext = %v", ciphertext)
		debug("tag        = %v", tag)
		debug("aad        = %v", aad)
		plaintext, err = cipher.decrypt(cek, iv, ciphertext, tag, aad)
		if err == nil {
			break
		}
	}

	if plaintext == nil {
		return nil, errors.New("failed to decrypt key")
	}

	if h.Compression == jwa.Deflate {
		output := bytes.Buffer{}
		w, _ := flate.NewWriter(&output, 1)
		in := plaintext
		for len(in) > 0 {
			n, err := w.Write(in)
			if err != nil {
				return nil, err
			}
			in = in[n:]
		}
		if err := w.Close(); err != nil {
			return nil, err
		}
		plaintext = output.Bytes()
	}

	return plaintext, nil
}
<|endoftext|>


<commit_before>package jws

import (
	"bytes"
	"encoding/json"
	"net/http"
	"strings"

	"github.com/SermoDigital/jose"
	"github.com/SermoDigital/jose/crypto"
)

// JWS implements a JWS per RFC 7515.
type JWS interface {
	// Payload Returns the payload.
	Payload() interface{}

	// SetPayload sets the payload with the given value.
	SetPayload(interface{})

	// Protected returns the JWS' Protected Header.
	// i represents the index of the Protected Header.
	// Left empty, it defaults to 0.
	Protected(...int) jose.Protected

	// Header returns the JWS' unprotected Header.
	// i represents the index of the Protected Header.
	// Left empty, it defaults to 0.
	Header(...int) jose.Header

	// Verify validates the current JWS' signature as-is. Refer to
	// ValidateMulti for more information.
	Verify(key interface{}, method crypto.SigningMethod) error

	// ValidateMulti validates the current JWS' signature as-is. Since it's
	// meant to be called after parsing a stream of bytes into a JWS, it
	// shouldn't do any internal parsing like the Sign, Flat, Compact, or
	// General methods do.
	VerifyMulti(keys []interface{}, methods []crypto.SigningMethod, o *SigningOpts) error

	// VerifyCallback validates the current JWS' signature as-is. It
	// accepts a callback function that can be used to access header
	// parameters to lookup needed information. For example, looking
	// up the "kid" parameter.
	// The return slice must be a slice of keys used in the verification
	// of the JWS.
	VerifyCallback(fn VerifyCallback, methods []crypto.SigningMethod, o *SigningOpts) error

	// General serializes the JWS into its "general" form per
	// https://tools.ietf.org/html/rfc7515#section-7.2.1
	General(keys ...interface{}) ([]byte, error)

	// Flat serializes the JWS to its "flattened" form per
	// https://tools.ietf.org/html/rfc7515#section-7.2.2
	Flat(key interface{}) ([]byte, error)

	// Compact serializes the JWS into its "compact" form per
	// https://tools.ietf.org/html/rfc7515#section-7.1
	Compact(key interface{}) ([]byte, error)

	// IsJWT returns true if the JWS is a JWT.
	IsJWT() bool
}

// jws represents a specific jws.
type jws struct {
	payload *payload
	plcache rawBase64
	clean   bool

	sb []sigHead

	isJWT bool
}

// Payload returns the jws' payload.
func (j *jws) Payload() interface{} { return j.payload.v }

// SetPayload sets the jws' raw, unexported payload.
func (j *jws) SetPayload(val interface{}) { j.payload.v = val }

// Protected returns the JWS' Protected Header.
// i represents the index of the Protected Header.
// Left empty, it defaults to 0.
func (j *jws) Protected(i ...int) jose.Protected {
	if len(i) == 0 {
		return j.sb[0].protected
	}
	return j.sb[i[0]].protected
}

// Header returns the JWS' unprotected Header.
// i represents the index of the Protected Header.
// Left empty, it defaults to 0.
func (j *jws) Header(i ...int) jose.Header {
	if len(i) == 0 {
		return j.sb[0].unprotected
	}
	return j.sb[i[0]].unprotected
}

// sigHead represents the 'signatures' member of the jws' "general"
// serialization form per
// https://tools.ietf.org/html/rfc7515#section-7.2.1
//
// It's embedded inside the "flat" structure in order to properly
// create the "flat" jws.
type sigHead struct {
	Protected   rawBase64        `json:"protected,omitempty"`
	Unprotected rawBase64        `json:"header,omitempty"`
	Signature   crypto.Signature `json:"signature"`

	protected   jose.Protected
	unprotected jose.Header
	clean       bool

	method crypto.SigningMethod
}

func (s *sigHead) unmarshal() error {
	if err := s.protected.UnmarshalJSON(s.Protected); err != nil {
		return err
	}
	if err := s.unprotected.UnmarshalJSON(s.Unprotected); err != nil {
		return err
	}
	return nil
}

// New creates a JWS with the provided crypto.SigningMethods.
func New(content interface{}, methods ...crypto.SigningMethod) JWS {
	sb := make([]sigHead, len(methods))
	for i := range methods {
		sb[i] = sigHead{
			protected: jose.Protected{
				"alg": methods[i].Alg(),
			},
			unprotected: jose.Header{},
			method:      methods[i],
		}
	}
	return &jws{
		payload: &payload{v: content},
		sb:      sb,
	}
}

func (s *sigHead) assignMethod(p jose.Protected) error {
	alg, ok := p.Get("alg").(string)
	if !ok {
		return ErrNoAlgorithm
	}

	sm := GetSigningMethod(alg)
	if sm == nil {
		return ErrNoAlgorithm
	}

	s.method = sm
	return nil
}

type generic struct {
	Payload rawBase64 `json:"payload"`
	sigHead
	Signatures []sigHead `json:"signatures,omitempty"`
}

// Parse parses any of the three serialized jws forms into a physical
// jws per https://tools.ietf.org/html/rfc7515#section-5.2
//
// It accepts a json.Unmarshaler in order to properly parse
// the payload. In order to keep the caller from having to do extra
// parsing of the payload, a json.Unmarshaler can be passed
// which will be then to unmarshal the payload however the caller
// wishes. Do note that if json.Unmarshal returns an error the
// original payload will be used as if no json.Unmarshaler was
// passed.
//
// Internally, Parse applies some heuristics and then calls either
// ParseGeneral, ParseFlat, or ParseCompact.
// It should only be called if, for whatever reason, you do not
// know which form the serialized JWT is in.
//
// It cannot parse a JWT.
func Parse(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	// Try and unmarshal into a generic struct that'll
	// hopefully hold either of the two JSON serialization
	// formats.
	var g generic

	// Not valid JSON. Let's try compact.
	if err := json.Unmarshal(encoded, &g); err != nil {
		return ParseCompact(encoded, u...)
	}

	if g.Signatures == nil {
		return g.parseFlat(u...)
	}
	return g.parseGeneral(u...)
}

// ParseGeneral parses a jws serialized into its "general" form per
// https://tools.ietf.org/html/rfc7515#section-7.2.1
// into a physical jws per
// https://tools.ietf.org/html/rfc7515#section-5.2
//
// For information on the json.Unmarshaler parameter, see Parse.
func ParseGeneral(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	var g generic
	if err := json.Unmarshal(encoded, &g); err != nil {
		return nil, err
	}
	return g.parseGeneral(u...)
}

func (g *generic) parseGeneral(u ...json.Unmarshaler) (JWS, error) {

	var p payload
	if len(u) > 0 {
		p.u = u[0]
	}

	if err := p.UnmarshalJSON(g.Payload); err != nil {
		return nil, err
	}

	for i := range g.Signatures {
		if err := g.Signatures[i].unmarshal(); err != nil {
			return nil, err
		}
		if err := checkHeaders(jose.Header(g.Signatures[i].protected), g.Signatures[i].unprotected); err != nil {
			return nil, err
		}

		if err := g.Signatures[i].assignMethod(g.Signatures[i].protected); err != nil {
			return nil, err
		}

		g.clean = true
	}

	return &jws{
		payload: &p,
		plcache: g.Payload,
		clean:   true,
		sb:      g.Signatures,
	}, nil
}

// ParseFlat parses a jws serialized into its "flat" form per
// https://tools.ietf.org/html/rfc7515#section-7.2.2
// into a physical jws per
// https://tools.ietf.org/html/rfc7515#section-5.2
//
// For information on the json.Unmarshaler parameter, see Parse.
func ParseFlat(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	var g generic
	if err := json.Unmarshal(encoded, &g); err != nil {
		return nil, err
	}
	return g.parseFlat(u...)
}

func (g *generic) parseFlat(u ...json.Unmarshaler) (JWS, error) {

	var p payload
	if len(u) > 0 {
		p.u = u[0]
	}

	if err := p.UnmarshalJSON(g.Payload); err != nil {
		return nil, err
	}

	if err := g.sigHead.unmarshal(); err != nil {
		return nil, err
	}
	g.sigHead.clean = true

	if err := checkHeaders(jose.Header(g.sigHead.protected), g.sigHead.unprotected); err != nil {
		return nil, err
	}

	if err := g.sigHead.assignMethod(g.sigHead.protected); err != nil {
		return nil, err
	}

	return &jws{
		payload: &p,
		plcache: g.Payload,
		clean:   true,
		sb:      []sigHead{g.sigHead},
	}, nil
}

// ParseCompact parses a jws serialized into its "compact" form per
// https://tools.ietf.org/html/rfc7515#section-7.1
// into a physical jws per
// https://tools.ietf.org/html/rfc7515#section-5.2
//
// For information on the json.Unmarshaler parameter, see Parse.
func ParseCompact(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	return parseCompact(encoded, false)
}

func parseCompact(encoded []byte, jwt bool, u ...json.Unmarshaler) (*jws, error) {

	// This section loosely follows
	// https://tools.ietf.org/html/rfc7519#section-7.2
	// because it's used to parse _both_ jws and JWTs.

	parts := bytes.Split(encoded, []byte{'.'})
	if len(parts) != 3 {
		return nil, ErrNotCompact
	}

	var p jose.Protected
	if err := p.UnmarshalJSON(parts[0]); err != nil {
		return nil, err
	}

	s := sigHead{
		Protected: parts[0],
		protected: p,
		Signature: parts[2],
		clean:     true,
	}

	if err := s.assignMethod(p); err != nil {
		return nil, err
	}

	var pl payload
	if len(u) > 0 {
		pl.u = u[0]
	}

	j := jws{
		payload: &pl,
		plcache: parts[1],
		sb:      []sigHead{s},
		isJWT:   jwt,
	}

	if err := j.payload.UnmarshalJSON(parts[1]); err != nil {
		return nil, err
	}

	j.clean = true

	if err := j.sb[0].Signature.UnmarshalJSON(parts[2]); err != nil {
		return nil, err
	}

	// https://tools.ietf.org/html/rfc7519#section-7.2.8
	cty, ok := p.Get("cty").(string)
	if ok && cty == "JWT" {
		return &j, ErrHoldsJWE
	}
	return &j, nil
}

var (
	// JWSFormKey is the form "key" which should be used inside
	// ParseFromRequest if the request is a multipart.Form.
	JWSFormKey = "access_token"

	// MaxMemory is maximum amount of memory which should be used
	// inside ParseFromRequest while parsing the multipart.Form
	// if the request is a multipart.Form.
	MaxMemory int64 = 10e6
)

// Format specifies which "format" the JWS is in -- Flat, General,
// or compact. Additionally, constants for JWT/Unknown are added.
type Format uint8

const (
	// Unknown format.
	Unknown Format = iota

	// Flat format.
	Flat

	// General format.
	General

	// Compact format.
	Compact
)

var parseJumpTable = [^Format(0)]func([]byte, ...json.Unmarshaler) (JWS, error){
	Unknown: Parse,
	Flat:    ParseFlat,
	General: ParseGeneral,
	Compact: ParseCompact,
}

func init() {
	for i := range parseJumpTable {
		if parseJumpTable[i] == nil {
			parseJumpTable[i] = Parse
		}
	}
}

func fromHeader(req *http.Request) ([]byte, bool) {
	if ah := req.Header.Get("Authorization"); ah != "" && len(ah) > 6 && strings.EqualFold(ah[0:6], "BEARER") {
		return []byte(ah[:7]), true
	}
	return nil, false
}

func fromForm(req *http.Request) ([]byte, bool) {
	if err := req.ParseMultipartForm(MaxMemory); err != nil {
		return nil, false
	}
	if tokStr := req.Form.Get(JWSFormKey); tokStr != "" {
		return []byte(tokStr), true
	}
	return nil, false
}

// ParseFromRequest tries to find the JWS in an http.Request.
// This method will call ParseMultipartForm if there's no token in the header.
func ParseFromRequest(req *http.Request, format Format, u ...json.Unmarshaler) (JWS, error) {
	if b, ok := fromHeader(req); ok {
		return parseJumpTable[format](b, u...)
	}
	if b, ok := fromForm(req); ok {
		return parseJumpTable[format](b, u...)
	}
	return nil, ErrNoTokenInRequest
}

// IgnoreDupes should be set to true if the internal duplicate header key check
// should ignore duplicate Header keys instead of reporting an error when
// duplicate Header keys are found.
//
// Note:
//     Duplicate Header keys are defined in
//     https://tools.ietf.org/html/rfc7515#section-5.2
//     meaning keys that both the protected and unprotected
//     Headers possess.
var IgnoreDupes bool

// checkHeaders returns an error per the constraints described in
// IgnoreDupes' comment.
func checkHeaders(a, b jose.Header) error {
	if len(a)+len(b) == 0 {
		return ErrTwoEmptyHeaders
	}
	for key := range a {
		if b.Has(key) && !IgnoreDupes {
			return ErrDuplicateHeaderParameter
		}
	}
	return nil
}

var _ JWS = (*jws)(nil)
<commit_msg>fromHeader should return strings behind "BEARER"<commit_after>package jws

import (
	"bytes"
	"encoding/json"
	"net/http"
	"strings"

	"github.com/SermoDigital/jose"
	"github.com/SermoDigital/jose/crypto"
)

// JWS implements a JWS per RFC 7515.
type JWS interface {
	// Payload Returns the payload.
	Payload() interface{}

	// SetPayload sets the payload with the given value.
	SetPayload(interface{})

	// Protected returns the JWS' Protected Header.
	// i represents the index of the Protected Header.
	// Left empty, it defaults to 0.
	Protected(...int) jose.Protected

	// Header returns the JWS' unprotected Header.
	// i represents the index of the Protected Header.
	// Left empty, it defaults to 0.
	Header(...int) jose.Header

	// Verify validates the current JWS' signature as-is. Refer to
	// ValidateMulti for more information.
	Verify(key interface{}, method crypto.SigningMethod) error

	// ValidateMulti validates the current JWS' signature as-is. Since it's
	// meant to be called after parsing a stream of bytes into a JWS, it
	// shouldn't do any internal parsing like the Sign, Flat, Compact, or
	// General methods do.
	VerifyMulti(keys []interface{}, methods []crypto.SigningMethod, o *SigningOpts) error

	// VerifyCallback validates the current JWS' signature as-is. It
	// accepts a callback function that can be used to access header
	// parameters to lookup needed information. For example, looking
	// up the "kid" parameter.
	// The return slice must be a slice of keys used in the verification
	// of the JWS.
	VerifyCallback(fn VerifyCallback, methods []crypto.SigningMethod, o *SigningOpts) error

	// General serializes the JWS into its "general" form per
	// https://tools.ietf.org/html/rfc7515#section-7.2.1
	General(keys ...interface{}) ([]byte, error)

	// Flat serializes the JWS to its "flattened" form per
	// https://tools.ietf.org/html/rfc7515#section-7.2.2
	Flat(key interface{}) ([]byte, error)

	// Compact serializes the JWS into its "compact" form per
	// https://tools.ietf.org/html/rfc7515#section-7.1
	Compact(key interface{}) ([]byte, error)

	// IsJWT returns true if the JWS is a JWT.
	IsJWT() bool
}

// jws represents a specific jws.
type jws struct {
	payload *payload
	plcache rawBase64
	clean   bool

	sb []sigHead

	isJWT bool
}

// Payload returns the jws' payload.
func (j *jws) Payload() interface{} { return j.payload.v }

// SetPayload sets the jws' raw, unexported payload.
func (j *jws) SetPayload(val interface{}) { j.payload.v = val }

// Protected returns the JWS' Protected Header.
// i represents the index of the Protected Header.
// Left empty, it defaults to 0.
func (j *jws) Protected(i ...int) jose.Protected {
	if len(i) == 0 {
		return j.sb[0].protected
	}
	return j.sb[i[0]].protected
}

// Header returns the JWS' unprotected Header.
// i represents the index of the Protected Header.
// Left empty, it defaults to 0.
func (j *jws) Header(i ...int) jose.Header {
	if len(i) == 0 {
		return j.sb[0].unprotected
	}
	return j.sb[i[0]].unprotected
}

// sigHead represents the 'signatures' member of the jws' "general"
// serialization form per
// https://tools.ietf.org/html/rfc7515#section-7.2.1
//
// It's embedded inside the "flat" structure in order to properly
// create the "flat" jws.
type sigHead struct {
	Protected   rawBase64        `json:"protected,omitempty"`
	Unprotected rawBase64        `json:"header,omitempty"`
	Signature   crypto.Signature `json:"signature"`

	protected   jose.Protected
	unprotected jose.Header
	clean       bool

	method crypto.SigningMethod
}

func (s *sigHead) unmarshal() error {
	if err := s.protected.UnmarshalJSON(s.Protected); err != nil {
		return err
	}
	if err := s.unprotected.UnmarshalJSON(s.Unprotected); err != nil {
		return err
	}
	return nil
}

// New creates a JWS with the provided crypto.SigningMethods.
func New(content interface{}, methods ...crypto.SigningMethod) JWS {
	sb := make([]sigHead, len(methods))
	for i := range methods {
		sb[i] = sigHead{
			protected: jose.Protected{
				"alg": methods[i].Alg(),
			},
			unprotected: jose.Header{},
			method:      methods[i],
		}
	}
	return &jws{
		payload: &payload{v: content},
		sb:      sb,
	}
}

func (s *sigHead) assignMethod(p jose.Protected) error {
	alg, ok := p.Get("alg").(string)
	if !ok {
		return ErrNoAlgorithm
	}

	sm := GetSigningMethod(alg)
	if sm == nil {
		return ErrNoAlgorithm
	}

	s.method = sm
	return nil
}

type generic struct {
	Payload rawBase64 `json:"payload"`
	sigHead
	Signatures []sigHead `json:"signatures,omitempty"`
}

// Parse parses any of the three serialized jws forms into a physical
// jws per https://tools.ietf.org/html/rfc7515#section-5.2
//
// It accepts a json.Unmarshaler in order to properly parse
// the payload. In order to keep the caller from having to do extra
// parsing of the payload, a json.Unmarshaler can be passed
// which will be then to unmarshal the payload however the caller
// wishes. Do note that if json.Unmarshal returns an error the
// original payload will be used as if no json.Unmarshaler was
// passed.
//
// Internally, Parse applies some heuristics and then calls either
// ParseGeneral, ParseFlat, or ParseCompact.
// It should only be called if, for whatever reason, you do not
// know which form the serialized JWT is in.
//
// It cannot parse a JWT.
func Parse(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	// Try and unmarshal into a generic struct that'll
	// hopefully hold either of the two JSON serialization
	// formats.
	var g generic

	// Not valid JSON. Let's try compact.
	if err := json.Unmarshal(encoded, &g); err != nil {
		return ParseCompact(encoded, u...)
	}

	if g.Signatures == nil {
		return g.parseFlat(u...)
	}
	return g.parseGeneral(u...)
}

// ParseGeneral parses a jws serialized into its "general" form per
// https://tools.ietf.org/html/rfc7515#section-7.2.1
// into a physical jws per
// https://tools.ietf.org/html/rfc7515#section-5.2
//
// For information on the json.Unmarshaler parameter, see Parse.
func ParseGeneral(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	var g generic
	if err := json.Unmarshal(encoded, &g); err != nil {
		return nil, err
	}
	return g.parseGeneral(u...)
}

func (g *generic) parseGeneral(u ...json.Unmarshaler) (JWS, error) {

	var p payload
	if len(u) > 0 {
		p.u = u[0]
	}

	if err := p.UnmarshalJSON(g.Payload); err != nil {
		return nil, err
	}

	for i := range g.Signatures {
		if err := g.Signatures[i].unmarshal(); err != nil {
			return nil, err
		}
		if err := checkHeaders(jose.Header(g.Signatures[i].protected), g.Signatures[i].unprotected); err != nil {
			return nil, err
		}

		if err := g.Signatures[i].assignMethod(g.Signatures[i].protected); err != nil {
			return nil, err
		}

		g.clean = true
	}

	return &jws{
		payload: &p,
		plcache: g.Payload,
		clean:   true,
		sb:      g.Signatures,
	}, nil
}

// ParseFlat parses a jws serialized into its "flat" form per
// https://tools.ietf.org/html/rfc7515#section-7.2.2
// into a physical jws per
// https://tools.ietf.org/html/rfc7515#section-5.2
//
// For information on the json.Unmarshaler parameter, see Parse.
func ParseFlat(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	var g generic
	if err := json.Unmarshal(encoded, &g); err != nil {
		return nil, err
	}
	return g.parseFlat(u...)
}

func (g *generic) parseFlat(u ...json.Unmarshaler) (JWS, error) {

	var p payload
	if len(u) > 0 {
		p.u = u[0]
	}

	if err := p.UnmarshalJSON(g.Payload); err != nil {
		return nil, err
	}

	if err := g.sigHead.unmarshal(); err != nil {
		return nil, err
	}
	g.sigHead.clean = true

	if err := checkHeaders(jose.Header(g.sigHead.protected), g.sigHead.unprotected); err != nil {
		return nil, err
	}

	if err := g.sigHead.assignMethod(g.sigHead.protected); err != nil {
		return nil, err
	}

	return &jws{
		payload: &p,
		plcache: g.Payload,
		clean:   true,
		sb:      []sigHead{g.sigHead},
	}, nil
}

// ParseCompact parses a jws serialized into its "compact" form per
// https://tools.ietf.org/html/rfc7515#section-7.1
// into a physical jws per
// https://tools.ietf.org/html/rfc7515#section-5.2
//
// For information on the json.Unmarshaler parameter, see Parse.
func ParseCompact(encoded []byte, u ...json.Unmarshaler) (JWS, error) {
	return parseCompact(encoded, false)
}

func parseCompact(encoded []byte, jwt bool, u ...json.Unmarshaler) (*jws, error) {

	// This section loosely follows
	// https://tools.ietf.org/html/rfc7519#section-7.2
	// because it's used to parse _both_ jws and JWTs.

	parts := bytes.Split(encoded, []byte{'.'})
	if len(parts) != 3 {
		return nil, ErrNotCompact
	}

	var p jose.Protected
	if err := p.UnmarshalJSON(parts[0]); err != nil {
		return nil, err
	}

	s := sigHead{
		Protected: parts[0],
		protected: p,
		Signature: parts[2],
		clean:     true,
	}

	if err := s.assignMethod(p); err != nil {
		return nil, err
	}

	var pl payload
	if len(u) > 0 {
		pl.u = u[0]
	}

	j := jws{
		payload: &pl,
		plcache: parts[1],
		sb:      []sigHead{s},
		isJWT:   jwt,
	}

	if err := j.payload.UnmarshalJSON(parts[1]); err != nil {
		return nil, err
	}

	j.clean = true

	if err := j.sb[0].Signature.UnmarshalJSON(parts[2]); err != nil {
		return nil, err
	}

	// https://tools.ietf.org/html/rfc7519#section-7.2.8
	cty, ok := p.Get("cty").(string)
	if ok && cty == "JWT" {
		return &j, ErrHoldsJWE
	}
	return &j, nil
}

var (
	// JWSFormKey is the form "key" which should be used inside
	// ParseFromRequest if the request is a multipart.Form.
	JWSFormKey = "access_token"

	// MaxMemory is maximum amount of memory which should be used
	// inside ParseFromRequest while parsing the multipart.Form
	// if the request is a multipart.Form.
	MaxMemory int64 = 10e6
)

// Format specifies which "format" the JWS is in -- Flat, General,
// or compact. Additionally, constants for JWT/Unknown are added.
type Format uint8

const (
	// Unknown format.
	Unknown Format = iota

	// Flat format.
	Flat

	// General format.
	General

	// Compact format.
	Compact
)

var parseJumpTable = [^Format(0)]func([]byte, ...json.Unmarshaler) (JWS, error){
	Unknown: Parse,
	Flat:    ParseFlat,
	General: ParseGeneral,
	Compact: ParseCompact,
}

func init() {
	for i := range parseJumpTable {
		if parseJumpTable[i] == nil {
			parseJumpTable[i] = Parse
		}
	}
}

func fromHeader(req *http.Request) ([]byte, bool) {
	if ah := req.Header.Get("Authorization"); ah != "" && len(ah) > 6 && strings.EqualFold(ah[0:6], "BEARER") {
		return []byte(ah[7:]), true
	}
	return nil, false
}

func fromForm(req *http.Request) ([]byte, bool) {
	if err := req.ParseMultipartForm(MaxMemory); err != nil {
		return nil, false
	}
	if tokStr := req.Form.Get(JWSFormKey); tokStr != "" {
		return []byte(tokStr), true
	}
	return nil, false
}

// ParseFromRequest tries to find the JWS in an http.Request.
// This method will call ParseMultipartForm if there's no token in the header.
func ParseFromRequest(req *http.Request, format Format, u ...json.Unmarshaler) (JWS, error) {
	if b, ok := fromHeader(req); ok {
		return parseJumpTable[format](b, u...)
	}
	if b, ok := fromForm(req); ok {
		return parseJumpTable[format](b, u...)
	}
	return nil, ErrNoTokenInRequest
}

// IgnoreDupes should be set to true if the internal duplicate header key check
// should ignore duplicate Header keys instead of reporting an error when
// duplicate Header keys are found.
//
// Note:
//     Duplicate Header keys are defined in
//     https://tools.ietf.org/html/rfc7515#section-5.2
//     meaning keys that both the protected and unprotected
//     Headers possess.
var IgnoreDupes bool

// checkHeaders returns an error per the constraints described in
// IgnoreDupes' comment.
func checkHeaders(a, b jose.Header) error {
	if len(a)+len(b) == 0 {
		return ErrTwoEmptyHeaders
	}
	for key := range a {
		if b.Has(key) && !IgnoreDupes {
			return ErrDuplicateHeaderParameter
		}
	}
	return nil
}

var _ JWS = (*jws)(nil)
<|endoftext|>


<commit_before>// Package kinesis implements a batch producer built on top of the official AWS SDK.
package kinesis

import (
	"errors"
	"time"

	"github.com/apex/log"
	k "github.com/aws/aws-sdk-go/service/kinesis"
)

// Size limits as defined by http://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecords.html.
const (
	maxRecordSize  = 1 << 20 // 1MiB
	maxRequestSize = 5 << 20 // 5MiB
)

// Errors.
var (
	ErrRecordSizeExceeded = errors.New("kinesis: record size exceeded")
)

// Producer batches records.
type Producer struct {
	Config
	records chan *k.PutRecordsRequestEntry
	done    chan struct{}
}

// New producer with the given config.
func New(config Config) *Producer {
	config.defaults()
	return &Producer{
		Config:  config,
		records: make(chan *k.PutRecordsRequestEntry, config.BacklogSize),
		done:    make(chan struct{}),
	}
}

// Put record `data` using `partitionKey`. This method is thread-safe.
func (p *Producer) Put(data []byte, partitionKey string) error {
	if len(data) > maxRecordSize {
		return ErrRecordSizeExceeded
	}

	p.records <- &k.PutRecordsRequestEntry{
		Data:         data,
		PartitionKey: &partitionKey,
	}

	return nil
}

// Start the producer.
func (p *Producer) Start() {
	go p.loop()
}

// Stop the producer. Flushes any in-flight data.
func (p *Producer) Stop() {
	p.Logger.WithField("backlog", len(p.records)).Info("stopping producer")

	// drain
	p.done <- struct{}{}
	close(p.records)

	// wait
	<-p.done

	p.Logger.Info("stopped producer")
}

// loop and flush at the configured interval, or when the buffer is exceeded.
func (p *Producer) loop() {
	buf := make([]*k.PutRecordsRequestEntry, 0, p.BufferSize)
	tick := time.NewTicker(p.FlushInterval)
	drain := false

	defer tick.Stop()
	defer close(p.done)

	for {
		select {
		case record := <-p.records:
			buf = append(buf, record)

			if len(buf) >= p.BufferSize {
				p.flush(buf, "buffer size")
				buf = nil
			}

			if drain && len(p.records) == 0 {
				p.Logger.Info("drained")
				return
			}
		case <-tick.C:
			if len(buf) > 0 {
				p.flush(buf, "interval")
				buf = nil
			}
		case <-p.done:
			drain = true

			if len(p.records) == 0 {
				return
			}
		}
	}
}

// flush records and retry failures if necessary.
func (p *Producer) flush(records []*k.PutRecordsRequestEntry, reason string) {
	p.Logger.WithFields(log.Fields{
		"records": len(records),
		"reason":  reason,
	}).Info("flush")

	out, err := p.Client.PutRecords(&k.PutRecordsInput{
		StreamName: &p.StreamName,
		Records:    records,
	})

	if err != nil {
		p.Logger.WithError(err).Error("flush")
		p.backoff(len(records))
		p.flush(records, "error")
		return
	}

	failed := *out.FailedRecordCount
	if failed == 0 {
		p.Backoff.Reset()
		return
	}

	p.backoff(int(failed))
	p.flush(failures(records, out.Records), "retry")
}

// calculates backoff duration and pauses execution
func (p *Producer) backoff(failed int) {
	backoff := p.Backoff.Duration()

	p.Logger.WithFields(log.Fields{
		"failures": failed,
		"backoff":  backoff,
	}).Warn("put failures")

	time.Sleep(backoff)
}

// failures returns the failed records as indicated in the response.
func failures(records []*k.PutRecordsRequestEntry, response []*k.PutRecordsResultEntry) (out []*k.PutRecordsRequestEntry) {
	for i, record := range response {
		if record.ErrorCode != nil {
			out = append(out, records[i])
		}
	}
	return
}
<commit_msg>Flush the buffered records when Stop()'d<commit_after>// Package kinesis implements a batch producer built on top of the official AWS SDK.
package kinesis

import (
	"errors"
	"time"

	"github.com/apex/log"
	k "github.com/aws/aws-sdk-go/service/kinesis"
)

// Size limits as defined by http://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecords.html.
const (
	maxRecordSize  = 1 << 20 // 1MiB
	maxRequestSize = 5 << 20 // 5MiB
)

// Errors.
var (
	ErrRecordSizeExceeded = errors.New("kinesis: record size exceeded")
)

// Producer batches records.
type Producer struct {
	Config
	records chan *k.PutRecordsRequestEntry
	done    chan struct{}
}

// New producer with the given config.
func New(config Config) *Producer {
	config.defaults()
	return &Producer{
		Config:  config,
		records: make(chan *k.PutRecordsRequestEntry, config.BacklogSize),
		done:    make(chan struct{}),
	}
}

// Put record `data` using `partitionKey`. This method is thread-safe.
func (p *Producer) Put(data []byte, partitionKey string) error {
	if len(data) > maxRecordSize {
		return ErrRecordSizeExceeded
	}

	p.records <- &k.PutRecordsRequestEntry{
		Data:         data,
		PartitionKey: &partitionKey,
	}

	return nil
}

// Start the producer.
func (p *Producer) Start() {
	go p.loop()
}

// Stop the producer. Flushes any in-flight data.
func (p *Producer) Stop() {
	p.Logger.WithField("backlog", len(p.records)).Info("stopping producer")

	// drain
	p.done <- struct{}{}
	close(p.records)

	// wait
	<-p.done

	p.Logger.Info("stopped producer")
}

// loop and flush at the configured interval, or when the buffer is exceeded.
func (p *Producer) loop() {
	buf := make([]*k.PutRecordsRequestEntry, 0, p.BufferSize)
	tick := time.NewTicker(p.FlushInterval)
	drain := false

	defer tick.Stop()
	defer close(p.done)

	for {
		select {
		case record := <-p.records:
			buf = append(buf, record)

			if len(buf) >= p.BufferSize {
				p.flush(buf, "buffer size")
				buf = nil
			}

			if drain && len(p.records) == 0 {
				p.Logger.Info("drained")
				return
			}
		case <-tick.C:
			if len(buf) > 0 {
				p.flush(buf, "interval")
				buf = nil
			}
		case <-p.done:
			drain = true
			if len(buf) > 0 {
				p.flush(buf, "stopping")
				buf = nil
			}

			if len(p.records) == 0 {
				return
			}
		}
	}
}

// flush records and retry failures if necessary.
func (p *Producer) flush(records []*k.PutRecordsRequestEntry, reason string) {
	p.Logger.WithFields(log.Fields{
		"records": len(records),
		"reason":  reason,
	}).Info("flush")

	out, err := p.Client.PutRecords(&k.PutRecordsInput{
		StreamName: &p.StreamName,
		Records:    records,
	})

	if err != nil {
		p.Logger.WithError(err).Error("flush")
		p.backoff(len(records))
		p.flush(records, "error")
		return
	}

	failed := *out.FailedRecordCount
	if failed == 0 {
		p.Backoff.Reset()
		return
	}

	p.backoff(int(failed))
	p.flush(failures(records, out.Records), "retry")
}

// calculates backoff duration and pauses execution
func (p *Producer) backoff(failed int) {
	backoff := p.Backoff.Duration()

	p.Logger.WithFields(log.Fields{
		"failures": failed,
		"backoff":  backoff,
	}).Warn("put failures")

	time.Sleep(backoff)
}

// failures returns the failed records as indicated in the response.
func failures(records []*k.PutRecordsRequestEntry, response []*k.PutRecordsResultEntry) (out []*k.PutRecordsRequestEntry) {
	for i, record := range response {
		if record.ErrorCode != nil {
			out = append(out, records[i])
		}
	}
	return
}
<|endoftext|>


<commit_before>// base36 supports Base36 encoding and decoding.
package base36

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"strconv"
	"strings"

	"github.com/grokify/base36"
	"github.com/grokify/gotilla/math/bigint"
)

// Md5Base36Length is the length for a MD5 Base36 string
const Md5Base36Length int = 25

// Encode36 returns an encoded string given a byte array.
func Encode36(ba []byte) string {
	return Encode36HexString(hex.EncodeToString(ba))
}

// Encode36String returns an encoded string given a string.
func Encode36String(s string) string {
	return Encode36([]byte(s))
}

// Encode36HexString returns an encoded string given a string.
func Encode36HexString(s16 string) string {
	i := bigint.HexToInt(s16)
	return strings.ToLower(base36.EncodeBigInt(i))
}

// Decode36 returns a decoded byte array given an encoded byte array.
func Decode36(b36 []byte) ([]byte, error) {
	return Decode36String(string(b36))
}

// Decode36String returns a decoded byte array given an encoded string.
func Decode36String(s36 string) ([]byte, error) {
	bi := base36.DecodeBigInt(s36)
	return hex.DecodeString(bigint.IntToHex(bi))
}

// Md5Base36 returns a Base36 encoded MD5 hash of a string.
func Md5Base36(s string) string {
	return fmt.Sprintf(`%0`+strconv.Itoa(Md5Base36Length)+`s`,
		Encode36HexString(fmt.Sprintf("%x", md5.Sum([]byte(s)))))
}

/*
GMP Versions

import "github.com/grokify/gmp"

// Encode36String returns an encoded string given a byte array.
func Encode36(ba []byte) string {
	s16 := hex.EncodeToString(ba)
	bi := gmp.NewInt(0)
	bi.SetString(s16, 16)
	return bi.InBase(36)
}

// Decode36String returns a decoded byte array given an encoded string.
func Decode36String(s36 string) ([]byte, error) {
	bi := gmp.NewInt(0)
	bi.SetString(s36, 36)
	s16 := bi.InBase(16)
	return hex.DecodeString(s16)
}

// Md5Base36 returns a Base36 encoded MD5 hash of a string.
func Md5Base36(s string) string {
	h := md5.New()
	io.WriteString(h, s)
	s16 := fmt.Sprintf("%x", h.Sum(nil))
	bi := gmp.NewInt(0)
	bi.SetString(s16, 16)
	return fmt.Sprintf("%025s", bi.InBase(36))
}
*/
<commit_msg>update style<commit_after>// base36 supports Base36 encoding and decoding.
package base36

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"strings"

	"github.com/grokify/base36"
	"github.com/grokify/gotilla/math/bigint"
)

// Md5Base36Length is the length for a MD5 Base36 string
const (
	Md5Base36Length int    = 25
	md5base36format string = `%025s`
)

// Encode36 returns an encoded string given a byte array.
func Encode36(ba []byte) string {
	return Encode36HexString(hex.EncodeToString(ba))
}

// Encode36String returns an encoded string given a string.
func Encode36String(s string) string {
	return Encode36([]byte(s))
}

// Encode36HexString returns an encoded string given a string.
func Encode36HexString(s16 string) string {
	i := bigint.HexToInt(s16)
	return strings.ToLower(base36.EncodeBigInt(i))
}

// Decode36 returns a decoded byte array given an encoded byte array.
func Decode36(b36 []byte) ([]byte, error) {
	return Decode36String(string(b36))
}

// Decode36String returns a decoded byte array given an encoded string.
func Decode36String(s36 string) ([]byte, error) {
	bi := base36.DecodeBigInt(s36)
	return hex.DecodeString(bigint.IntToHex(bi))
}

// Md5Base36 returns a Base36 encoded MD5 hash of a string.
func Md5Base36(s string) string {
	return fmt.Sprintf(md5base36format,
		Encode36HexString(fmt.Sprintf("%x", md5.Sum([]byte(s)))))
}

/*
GMP Versions

import "github.com/grokify/gmp"

// Encode36String returns an encoded string given a byte array.
func Encode36(ba []byte) string {
	s16 := hex.EncodeToString(ba)
	bi := gmp.NewInt(0)
	bi.SetString(s16, 16)
	return bi.InBase(36)
}

// Decode36String returns a decoded byte array given an encoded string.
func Decode36String(s36 string) ([]byte, error) {
	bi := gmp.NewInt(0)
	bi.SetString(s36, 36)
	s16 := bi.InBase(16)
	return hex.DecodeString(s16)
}

// Md5Base36 returns a Base36 encoded MD5 hash of a string.
func Md5Base36(s string) string {
	h := md5.New()
	io.WriteString(h, s)
	s16 := fmt.Sprintf("%x", h.Sum(nil))
	bi := gmp.NewInt(0)
	bi.SetString(s16, 16)
	return fmt.Sprintf("%025s", bi.InBase(36))
}
*/
<|endoftext|>


<commit_before>//  Copyright (c) 2013 Couchbase, 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.

package leveldb

import (
	"github.com/couchbaselabs/indexing/api"
	"github.com/jmhodges/levigo"
	"log"
)

func Create(name string) (*LevelDBEngine, error) {

	var ldb LevelDBEngine
	//FIXME move this to a dir
	ldb.name = name

	ldb.options = levigo.NewOptions()
	ldb.options.SetCreateIfMissing(true)
	ldb.options.SetErrorIfExists(true)

	//set filter policy
	filterPolicy := levigo.NewBloomFilter(10)
	ldb.options.SetFilterPolicy(filterPolicy)

	ldb.options.SetCompression(levigo.SnappyCompression)

	ldb.wo = levigo.NewWriteOptions()
	ldb.ro = levigo.NewReadOptions()

	var err error
	if ldb.c, err = levigo.Open(ldb.name, ldb.options); err != nil {
		return nil, err
	}

	//create a separate back-index
	if ldb.b, err = levigo.Open(ldb.name+"_back", ldb.options); err != nil {
		return nil, err
	}

	return &ldb, nil
}

func Open(name string) (*LevelDBEngine, error) {

	var ldb LevelDBEngine
	//FIXME move this to a dir
	ldb.name = name

	ldb.options = levigo.NewOptions()
	ldb.options.SetCreateIfMissing(false)

	//set filter policy
	filterPolicy := levigo.NewBloomFilter(10)
	ldb.options.SetFilterPolicy(filterPolicy)

	ldb.options.SetCompression(levigo.SnappyCompression)

	ldb.wo = levigo.NewWriteOptions()
	ldb.ro = levigo.NewReadOptions()

	var err error
	if ldb.c, err = levigo.Open(ldb.name, ldb.options); err != nil {
		return nil, err
	}

	//create a separate back-index
	if ldb.b, err = levigo.Open(ldb.name+"_back", ldb.options); err != nil {
		return nil, err
	}

	return &ldb, nil

}

func (ldb *LevelDBEngine) InsertMutation(k api.Key, v api.Value) error {

	var err error
	var backkey api.Key

	log.Printf("LevelDB Set Key - %s Value - %s", k.String(), v.String())

	//check if the docid exists in the back index
	if backkey, err = ldb.GetBackIndexEntry(v.Docid()); err != nil {
		log.Printf("Error locating backindex entry %v", err)
		return err
	} else if backkey.EncodedBytes() != nil {
		//there is already an entry in main index for this docid
		//delete from main index
		if err = ldb.c.Delete(ldb.wo, backkey.EncodedBytes()); err != nil {
			log.Printf("Error deleting entry from main index %v", err)
			return err
		}
	}
	//FIXME : Handle the case if old-value from backindex matches with the new-value(false mutation). Skip It.

	//set in main index
	if err = ldb.c.Put(ldb.wo, k.EncodedBytes(), v.EncodedBytes()); err != nil {
		return err
	}

	//set the back index entry <docid, encodedkey>
	if err = ldb.b.Put(ldb.wo, []byte(v.Docid()), k.EncodedBytes()); err != nil {
		return err
	}

	return err
}

func (ldb *LevelDBEngine) InsertMeta(metaid string, metavalue string) error {

	log.Printf("LevelDB Set Meta Key - %s, Value - %s", metaid, metavalue)
	var err error

	//meta values go to the back index
	if err = ldb.b.Put(ldb.wo, []byte(metaid), []byte(metavalue)); err != nil {
		return err
	}

	return err
}

func (ldb *LevelDBEngine) GetMeta(metaid string) (string, error) {

	var metavalue []byte
	var err error
	if metavalue, err = ldb.b.Get(ldb.ro, []byte(metaid)); err == nil {
		log.Printf("LevelDB Get Meta Key - %s, Value - %s", metaid, string(metavalue))
		return string(metavalue), nil
	}

	return "", err
}

func (ldb *LevelDBEngine) GetBackIndexEntry(docid string) (api.Key, error) {

	var k api.Key
	var kbyte []byte
	var err error

	log.Printf("LevelDB Get BackIndex Key - %s", docid)

	if kbyte, err = ldb.b.Get(ldb.ro, []byte(docid)); err != nil {
		return k, err
	}

	k, err = api.NewKeyFromEncodedBytes(kbyte)

	return k, err
}

func (ldb *LevelDBEngine) DeleteMutation(docid string) error {

	log.Printf("LevelDB Delete Key - %s", docid)
	var backkey api.Key
	var err error

	if backkey, err = ldb.GetBackIndexEntry(docid); err != nil {
		log.Printf("Error locating backindex entry %v", err)
		return err
	}

	//delete from main index
	if err = ldb.c.Delete(ldb.wo, backkey.EncodedBytes()); err != nil {
		log.Printf("Error deleting entry from main index %v", err)
		return err
	}

	//delete from the back index
	if err = ldb.b.Delete(ldb.wo, []byte(docid)); err != nil {
		log.Printf("Error deleting entry from back index %v", err)
		return err
	}

	return nil

}

func (ldb *LevelDBEngine) Close() error {
	//close the main index
	if ldb.c != nil {
		ldb.c.Close()
	}
	//close the back index
	if ldb.b != nil {
		ldb.b.Close()
	}
	return nil
}

func (ldb *LevelDBEngine) Destroy() error {
	var err error
	if err = ldb.Close(); err != nil {
		return err
	}
	//Destroy the main index
	if err = levigo.DestroyDatabase(ldb.name, ldb.options); err != nil {
		return err
	}
	//Destroy the back index
	if err = levigo.DestroyDatabase(ldb.name+"_back", ldb.options); err != nil {
		return err
	}
	return err
}

func (ldb *LevelDBEngine) GetKeySetForKeyRange(low api.Key, high api.Key,
	inclusion api.Inclusion, chkey chan api.Key, cherr chan error) {

	defer close(chkey)
	defer close(cherr)

	snap := ldb.c.NewSnapshot()
	defer ldb.c.ReleaseSnapshot(snap)

	ro := levigo.NewReadOptions()
	ro.SetSnapshot(snap)

	it := ldb.c.NewIterator(ro)
	defer it.Close()

	log.Printf("LevelDB Received Key Low - %s High - %s for Scan", low.String(), high.String())

	var lowkey []byte
	var err error

	if lowkey = low.EncodedBytes(); lowkey == nil {
		it.SeekToFirst()
	} else {
		it.Seek(lowkey)
	}

	var key api.Key
	for it = it; it.Valid(); it.Next() {
		if key, err = api.NewKeyFromEncodedBytes(it.Key()); err != nil {
			log.Printf("Error Converting from bytes %v to key %v. Skipping row", it.Key(), err)
			continue
		}

		log.Printf("LevelDB Got Key - %s", key.String())

		var highcmp int
		if high.EncodedBytes() == nil {
			highcmp = -1 //if high key is nil, iterate through the fullset
		} else {
			highcmp = key.Compare(high)
		}

		var lowcmp int
		if low.EncodedBytes() == nil {
			lowcmp = 1 //all keys are greater than nil
		} else {
			lowcmp = key.Compare(low)
		}

		if highcmp == 0 && (inclusion == api.Both || inclusion == api.High) {
			log.Printf("LevelDB Sending Key Equal to High Key")
			chkey <- key
		} else if lowcmp == 0 && (inclusion == api.Both || inclusion == api.Low) {
			log.Printf("LevelDB Sending Key Equal to Low Key")
			chkey <- key
		} else if (highcmp == -1) && (lowcmp == 1) { //key is between high and low
			if highcmp == -1 {
				log.Printf("LevelDB Sending Key Lesser Than High Key")
			} else if lowcmp == 1 {
				log.Printf("LevelDB Sending Key Greater Than Low Key")
			}
			chkey <- key
		} else {
			log.Printf("LevelDB not Sending Key")
			//if we have reached past the high key, no need to scan further
			if highcmp == 1 {
				break
			}
		}
	}

	//FIXME
	/*
	   if err := it.GetError() {
	       log.Printf("Error %v", err)
	       return err
	   }
	*/

}

func (ldb *LevelDBEngine) GetValueSetForKeyRange(low api.Key, high api.Key,
	inclusion api.Inclusion, chval chan api.Value, cherr chan error) {

	defer close(chval)
	defer close(cherr)

	snap := ldb.c.NewSnapshot()
	defer ldb.c.ReleaseSnapshot(snap)

	ro := levigo.NewReadOptions()
	ro.SetSnapshot(snap)

	it := ldb.c.NewIterator(ro)
	defer it.Close()

	log.Printf("LevelDB Received Key Low - %s High - %s Inclusion - %v for Scan", low.String(), high.String(), inclusion)

	var lowkey []byte
	var err error

	if lowkey = low.EncodedBytes(); lowkey == nil {
		it.SeekToFirst()
	} else {
		it.Seek(lowkey)
	}

	var key api.Key
	var val api.Value
	for it = it; it.Valid(); it.Next() {
		if key, err = api.NewKeyFromEncodedBytes(it.Key()); err != nil {
			log.Printf("Error Converting from bytes %v to key %v. Skipping row", it.Key(), err)
			continue
		}

		if val, err = api.NewValueFromEncodedBytes(it.Value()); err != nil {
			log.Printf("Error Converting from bytes %v to value %v, Skipping row", it.Value(), err)
			continue
		}

		log.Printf("LevelDB Got Value - %s", val.String())

		var highcmp int
		if high.EncodedBytes() == nil {
			highcmp = -1 //if high key is nil, iterate through the fullset
		} else {
			highcmp = key.Compare(high)
		}

		var lowcmp int
		if low.EncodedBytes() == nil {
			lowcmp = 1 //all keys are greater than nil
		} else {
			lowcmp = key.Compare(low)
		}

		if highcmp == 0 && (inclusion == api.Both || inclusion == api.High) {
			log.Printf("LevelDB Sending Value Equal to High Key")
			chval <- val
		} else if lowcmp == 0 && (inclusion == api.Both || inclusion == api.Low) {
			log.Printf("LevelDB Sending Value Equal to Low Key")
			chval <- val
		} else if (highcmp == -1) && (lowcmp == 1) { //key is between high and low
			if highcmp == -1 {
				log.Printf("LevelDB Sending Value Lesser Than High Key")
			} else if lowcmp == 1 {
				log.Printf("LevelDB Sending Value Greater Than Low Key")
			}
			chval <- val
		} else {
			log.Printf("LevelDB not Sending Value")
			//if we have reached past the high key, no need to scan further
			if highcmp == 1 {
				break
			}
		}
	}

	//FIXME
	/*
	   if err := it.GetError() {
	       log.Printf("Error %v", err)
	       return err
	   }
	*/

}

func (ldb *LevelDBEngine) CountRange(low api.Key, high api.Key, inclusion api.Inclusion) (
	uint64, error) {

	var count uint64

	snap := ldb.c.NewSnapshot()
	defer ldb.c.ReleaseSnapshot(snap)

	ro := levigo.NewReadOptions()
	ro.SetSnapshot(snap)

	it := ldb.c.NewIterator(ro)
	defer it.Close()

	log.Printf("LevelDB Received Key Low - %s High - %s for Scan", low.String(), high.String())

	var lowkey []byte
	var err error

	if lowkey = low.EncodedBytes(); lowkey == nil {
		it.SeekToFirst()
	} else {
		it.Seek(lowkey)
	}

	var key api.Key
	for it = it; it.Valid(); it.Next() {
		if key, err = api.NewKeyFromEncodedBytes(it.Key()); err != nil {
			log.Printf("Error Converting from bytes %v to key %v. Skipping row", it.Key(), err)
			continue
		}

		log.Printf("LevelDB Got Key - %s", key.String())

		var highcmp int
		if high.EncodedBytes() == nil {
			highcmp = -1 //if high key is nil, iterate through the fullset
		} else {
			highcmp = key.Compare(high)
		}

		var lowcmp int
		if low.EncodedBytes() == nil {
			lowcmp = 1 //all keys are greater than nil
		} else {
			lowcmp = key.Compare(low)
		}

		if highcmp == 0 && (inclusion == api.Both || inclusion == api.High) {
			log.Printf("LevelDB Sending Value Equal to High Key")
			count++
		} else if lowcmp == 0 && (inclusion == api.Both || inclusion == api.Low) {
			log.Printf("LevelDB Sending Value Equal to Low Key")
			count++
		} else if (highcmp == -1) && (lowcmp == 1) { //key is between high and low
			if highcmp == -1 {
				log.Printf("LevelDB Sending Value Lesser Than High Key")
			} else if lowcmp == 1 {
				log.Printf("LevelDB Sending Value Greater Than Low Key")
			}
			count++
		} else {
			log.Printf("LevelDB not Sending Value")
			//if we have reached past the high key, no need to scan further
			if highcmp == 1 {
				break
			}
		}
	}

	//FIXME
	/*
	   if err := it.GetError() {
	       log.Printf("Error %v", err)
	       return err
	   }
	*/

	return count, nil
}
<commit_msg>Set MaxOpenFiles for leveldb to 500<commit_after>//  Copyright (c) 2013 Couchbase, 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.

package leveldb

import (
	"github.com/couchbaselabs/indexing/api"
	"github.com/jmhodges/levigo"
	"log"
)

func Create(name string) (*LevelDBEngine, error) {

	var ldb LevelDBEngine
	//FIXME move this to a dir
	ldb.name = name

	ldb.options = levigo.NewOptions()
	ldb.options.SetCreateIfMissing(true)
	ldb.options.SetErrorIfExists(true)

	//set filter policy
	filterPolicy := levigo.NewBloomFilter(10)
	ldb.options.SetFilterPolicy(filterPolicy)

	ldb.options.SetCompression(levigo.SnappyCompression)
	ldb.options.SetMaxOpenFiles(500)

	ldb.wo = levigo.NewWriteOptions()
	ldb.ro = levigo.NewReadOptions()

	var err error
	if ldb.c, err = levigo.Open(ldb.name, ldb.options); err != nil {
		return nil, err
	}

	//create a separate back-index
	if ldb.b, err = levigo.Open(ldb.name+"_back", ldb.options); err != nil {
		return nil, err
	}

	return &ldb, nil
}

func Open(name string) (*LevelDBEngine, error) {

	var ldb LevelDBEngine
	//FIXME move this to a dir
	ldb.name = name

	ldb.options = levigo.NewOptions()
	ldb.options.SetCreateIfMissing(false)

	//set filter policy
	filterPolicy := levigo.NewBloomFilter(10)
	ldb.options.SetFilterPolicy(filterPolicy)

	ldb.options.SetCompression(levigo.SnappyCompression)
	ldb.options.SetMaxOpenFiles(500)

	ldb.wo = levigo.NewWriteOptions()
	ldb.ro = levigo.NewReadOptions()

	var err error
	if ldb.c, err = levigo.Open(ldb.name, ldb.options); err != nil {
		return nil, err
	}

	//create a separate back-index
	if ldb.b, err = levigo.Open(ldb.name+"_back", ldb.options); err != nil {
		return nil, err
	}

	return &ldb, nil

}

func (ldb *LevelDBEngine) InsertMutation(k api.Key, v api.Value) error {

	var err error
	var backkey api.Key

	log.Printf("LevelDB Set Key - %s Value - %s", k.String(), v.String())

	//check if the docid exists in the back index
	if backkey, err = ldb.GetBackIndexEntry(v.Docid()); err != nil {
		log.Printf("Error locating backindex entry %v", err)
		return err
	} else if backkey.EncodedBytes() != nil {
		//there is already an entry in main index for this docid
		//delete from main index
		if err = ldb.c.Delete(ldb.wo, backkey.EncodedBytes()); err != nil {
			log.Printf("Error deleting entry from main index %v", err)
			return err
		}
	}
	//FIXME : Handle the case if old-value from backindex matches with the new-value(false mutation). Skip It.

	//set in main index
	if err = ldb.c.Put(ldb.wo, k.EncodedBytes(), v.EncodedBytes()); err != nil {
		return err
	}

	//set the back index entry <docid, encodedkey>
	if err = ldb.b.Put(ldb.wo, []byte(v.Docid()), k.EncodedBytes()); err != nil {
		return err
	}

	return err
}

func (ldb *LevelDBEngine) InsertMeta(metaid string, metavalue string) error {

	log.Printf("LevelDB Set Meta Key - %s, Value - %s", metaid, metavalue)
	var err error

	//meta values go to the back index
	if err = ldb.b.Put(ldb.wo, []byte(metaid), []byte(metavalue)); err != nil {
		return err
	}

	return err
}

func (ldb *LevelDBEngine) GetMeta(metaid string) (string, error) {

	var metavalue []byte
	var err error
	if metavalue, err = ldb.b.Get(ldb.ro, []byte(metaid)); err == nil {
		log.Printf("LevelDB Get Meta Key - %s, Value - %s", metaid, string(metavalue))
		return string(metavalue), nil
	}

	return "", err
}

func (ldb *LevelDBEngine) GetBackIndexEntry(docid string) (api.Key, error) {

	var k api.Key
	var kbyte []byte
	var err error

	log.Printf("LevelDB Get BackIndex Key - %s", docid)

	if kbyte, err = ldb.b.Get(ldb.ro, []byte(docid)); err != nil {
		return k, err
	}

	k, err = api.NewKeyFromEncodedBytes(kbyte)

	return k, err
}

func (ldb *LevelDBEngine) DeleteMutation(docid string) error {

	log.Printf("LevelDB Delete Key - %s", docid)
	var backkey api.Key
	var err error

	if backkey, err = ldb.GetBackIndexEntry(docid); err != nil {
		log.Printf("Error locating backindex entry %v", err)
		return err
	}

	//delete from main index
	if err = ldb.c.Delete(ldb.wo, backkey.EncodedBytes()); err != nil {
		log.Printf("Error deleting entry from main index %v", err)
		return err
	}

	//delete from the back index
	if err = ldb.b.Delete(ldb.wo, []byte(docid)); err != nil {
		log.Printf("Error deleting entry from back index %v", err)
		return err
	}

	return nil

}

func (ldb *LevelDBEngine) Close() error {
	//close the main index
	if ldb.c != nil {
		ldb.c.Close()
	}
	//close the back index
	if ldb.b != nil {
		ldb.b.Close()
	}
	return nil
}

func (ldb *LevelDBEngine) Destroy() error {
	var err error
	if err = ldb.Close(); err != nil {
		return err
	}
	//Destroy the main index
	if err = levigo.DestroyDatabase(ldb.name, ldb.options); err != nil {
		return err
	}
	//Destroy the back index
	if err = levigo.DestroyDatabase(ldb.name+"_back", ldb.options); err != nil {
		return err
	}
	return err
}

func (ldb *LevelDBEngine) GetKeySetForKeyRange(low api.Key, high api.Key,
	inclusion api.Inclusion, chkey chan api.Key, cherr chan error) {

	defer close(chkey)
	defer close(cherr)

	snap := ldb.c.NewSnapshot()
	defer ldb.c.ReleaseSnapshot(snap)

	ro := levigo.NewReadOptions()
	ro.SetSnapshot(snap)

	it := ldb.c.NewIterator(ro)
	defer it.Close()

	log.Printf("LevelDB Received Key Low - %s High - %s for Scan", low.String(), high.String())

	var lowkey []byte
	var err error

	if lowkey = low.EncodedBytes(); lowkey == nil {
		it.SeekToFirst()
	} else {
		it.Seek(lowkey)
	}

	var key api.Key
	for it = it; it.Valid(); it.Next() {
		if key, err = api.NewKeyFromEncodedBytes(it.Key()); err != nil {
			log.Printf("Error Converting from bytes %v to key %v. Skipping row", it.Key(), err)
			continue
		}

		log.Printf("LevelDB Got Key - %s", key.String())

		var highcmp int
		if high.EncodedBytes() == nil {
			highcmp = -1 //if high key is nil, iterate through the fullset
		} else {
			highcmp = key.Compare(high)
		}

		var lowcmp int
		if low.EncodedBytes() == nil {
			lowcmp = 1 //all keys are greater than nil
		} else {
			lowcmp = key.Compare(low)
		}

		if highcmp == 0 && (inclusion == api.Both || inclusion == api.High) {
			log.Printf("LevelDB Sending Key Equal to High Key")
			chkey <- key
		} else if lowcmp == 0 && (inclusion == api.Both || inclusion == api.Low) {
			log.Printf("LevelDB Sending Key Equal to Low Key")
			chkey <- key
		} else if (highcmp == -1) && (lowcmp == 1) { //key is between high and low
			if highcmp == -1 {
				log.Printf("LevelDB Sending Key Lesser Than High Key")
			} else if lowcmp == 1 {
				log.Printf("LevelDB Sending Key Greater Than Low Key")
			}
			chkey <- key
		} else {
			log.Printf("LevelDB not Sending Key")
			//if we have reached past the high key, no need to scan further
			if highcmp == 1 {
				break
			}
		}
	}

	//FIXME
	/*
	   if err := it.GetError() {
	       log.Printf("Error %v", err)
	       return err
	   }
	*/

}

func (ldb *LevelDBEngine) GetValueSetForKeyRange(low api.Key, high api.Key,
	inclusion api.Inclusion, chval chan api.Value, cherr chan error) {

	defer close(chval)
	defer close(cherr)

	snap := ldb.c.NewSnapshot()
	defer ldb.c.ReleaseSnapshot(snap)

	ro := levigo.NewReadOptions()
	ro.SetSnapshot(snap)

	it := ldb.c.NewIterator(ro)
	defer it.Close()

	log.Printf("LevelDB Received Key Low - %s High - %s Inclusion - %v for Scan", low.String(), high.String(), inclusion)

	var lowkey []byte
	var err error

	if lowkey = low.EncodedBytes(); lowkey == nil {
		it.SeekToFirst()
	} else {
		it.Seek(lowkey)
	}

	var key api.Key
	var val api.Value
	for it = it; it.Valid(); it.Next() {
		if key, err = api.NewKeyFromEncodedBytes(it.Key()); err != nil {
			log.Printf("Error Converting from bytes %v to key %v. Skipping row", it.Key(), err)
			continue
		}

		if val, err = api.NewValueFromEncodedBytes(it.Value()); err != nil {
			log.Printf("Error Converting from bytes %v to value %v, Skipping row", it.Value(), err)
			continue
		}

		log.Printf("LevelDB Got Value - %s", val.String())

		var highcmp int
		if high.EncodedBytes() == nil {
			highcmp = -1 //if high key is nil, iterate through the fullset
		} else {
			highcmp = key.Compare(high)
		}

		var lowcmp int
		if low.EncodedBytes() == nil {
			lowcmp = 1 //all keys are greater than nil
		} else {
			lowcmp = key.Compare(low)
		}

		if highcmp == 0 && (inclusion == api.Both || inclusion == api.High) {
			log.Printf("LevelDB Sending Value Equal to High Key")
			chval <- val
		} else if lowcmp == 0 && (inclusion == api.Both || inclusion == api.Low) {
			log.Printf("LevelDB Sending Value Equal to Low Key")
			chval <- val
		} else if (highcmp == -1) && (lowcmp == 1) { //key is between high and low
			if highcmp == -1 {
				log.Printf("LevelDB Sending Value Lesser Than High Key")
			} else if lowcmp == 1 {
				log.Printf("LevelDB Sending Value Greater Than Low Key")
			}
			chval <- val
		} else {
			log.Printf("LevelDB not Sending Value")
			//if we have reached past the high key, no need to scan further
			if highcmp == 1 {
				break
			}
		}
	}

	//FIXME
	/*
	   if err := it.GetError() {
	       log.Printf("Error %v", err)
	       return err
	   }
	*/

}

func (ldb *LevelDBEngine) CountRange(low api.Key, high api.Key, inclusion api.Inclusion) (
	uint64, error) {

	var count uint64

	snap := ldb.c.NewSnapshot()
	defer ldb.c.ReleaseSnapshot(snap)

	ro := levigo.NewReadOptions()
	ro.SetSnapshot(snap)

	it := ldb.c.NewIterator(ro)
	defer it.Close()

	log.Printf("LevelDB Received Key Low - %s High - %s for Scan", low.String(), high.String())

	var lowkey []byte
	var err error

	if lowkey = low.EncodedBytes(); lowkey == nil {
		it.SeekToFirst()
	} else {
		it.Seek(lowkey)
	}

	var key api.Key
	for it = it; it.Valid(); it.Next() {
		if key, err = api.NewKeyFromEncodedBytes(it.Key()); err != nil {
			log.Printf("Error Converting from bytes %v to key %v. Skipping row", it.Key(), err)
			continue
		}

		log.Printf("LevelDB Got Key - %s", key.String())

		var highcmp int
		if high.EncodedBytes() == nil {
			highcmp = -1 //if high key is nil, iterate through the fullset
		} else {
			highcmp = key.Compare(high)
		}

		var lowcmp int
		if low.EncodedBytes() == nil {
			lowcmp = 1 //all keys are greater than nil
		} else {
			lowcmp = key.Compare(low)
		}

		if highcmp == 0 && (inclusion == api.Both || inclusion == api.High) {
			log.Printf("LevelDB Sending Value Equal to High Key")
			count++
		} else if lowcmp == 0 && (inclusion == api.Both || inclusion == api.Low) {
			log.Printf("LevelDB Sending Value Equal to Low Key")
			count++
		} else if (highcmp == -1) && (lowcmp == 1) { //key is between high and low
			if highcmp == -1 {
				log.Printf("LevelDB Sending Value Lesser Than High Key")
			} else if lowcmp == 1 {
				log.Printf("LevelDB Sending Value Greater Than Low Key")
			}
			count++
		} else {
			log.Printf("LevelDB not Sending Value")
			//if we have reached past the high key, no need to scan further
			if highcmp == 1 {
				break
			}
		}
	}

	//FIXME
	/*
	   if err := it.GetError() {
	       log.Printf("Error %v", err)
	       return err
	   }
	*/

	return count, nil
}
<|endoftext|>


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

package main

import (
	"flag"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"
)

var cmdDockerize = &Command{
	CustomFlags: true,
	UsageLine:   "dockerize",
	Short:       "Generates a Dockerfile for your Beego application",
	Long: `Dockerize generates a Dockerfile for your Beego Web Application.
  The Dockerfile will compile, get the dependencies with {{"godep"|bold}}, and set the entrypoint.
  `,
	PreRun: func(cmd *Command, args []string) { ShowShortVersionBanner() },
	Run:    dockerizeApp,
}

const dockerBuildTemplate = `FROM {{.BaseImage}}

# Godep for vendoring
RUN go get github.com/tools/godep

# Recompile the standard library without CGO
RUN CGO_ENABLED=0 go install -a std

ENV APP_DIR $GOPATH{{.Appdir}}
RUN mkdir -p $APP_DIR

# Set the entrypoint
ENTRYPOINT $APP_DIR/{{.Entrypoint}}
ADD . $APP_DIR

# Compile the binary and statically link
RUN cd $APP_DIR
RUN CGO_ENABLED=0 godep go build -ldflags '-d -w -s'

EXPOSE {{.Expose}}
`

type Dockerfile struct {
	BaseImage  string
	Appdir     string
	Entrypoint string
	Expose     string
}

var (
	expose    string
	baseImage string
)

func init() {
	fs := flag.NewFlagSet("dockerize", flag.ContinueOnError)
	fs.StringVar(&baseImage, "image", "library/golang", "Sets the base image of the Docker container.")
	fs.StringVar(&expose, "expose", "8080", "Port to expose in the Docker container.")
	cmdDockerize.Flag = *fs
}

func dockerizeApp(cmd *Command, args []string) int {
	cmd.Flag.Parse(args)

	logger.Info("Generating Dockerfile...")

	gopath := os.Getenv("GOPATH")
	dir, err := filepath.Abs(".")
	MustCheck(err)

	appdir := strings.Replace(dir, gopath, "", 1)

	// TODO: Check if the base image exists in Docker Hub
	_, entrypoint := path.Split(appdir)
	dockerfile := Dockerfile{
		BaseImage:  baseImage,
		Appdir:     appdir,
		Entrypoint: entrypoint,
		Expose:     expose,
	}

	generateDockerfile(dockerfile)
	return 0
}

func generateDockerfile(df Dockerfile) {
	t := template.Must(template.New("dockerBuildTemplate").Parse(dockerBuildTemplate)).Funcs(BeeFuncMap())

	f, err := os.Create("Dockerfile")
	if err != nil {
		logger.Fatalf("Error writing Dockerfile: %v", err.Error())
	}
	defer CloseFile(f)

	t.Execute(f, df)

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

package main

import (
	"flag"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"
)

var cmdDockerize = &Command{
	CustomFlags: true,
	UsageLine:   "dockerize",
	Short:       "Generates a Dockerfile for your Beego application",
	Long: `Dockerize generates a Dockerfile for your Beego Web Application.
  The Dockerfile will compile, get the dependencies with {{"godep"|bold}}, and set the entrypoint.

  {{"Example:"|bold}}
    $ bee dockerize -expose="3000,80,25"
  `,
	PreRun: func(cmd *Command, args []string) { ShowShortVersionBanner() },
	Run:    dockerizeApp,
}

const dockerBuildTemplate = `FROM {{.BaseImage}}

# Godep for vendoring
RUN go get github.com/tools/godep

# Recompile the standard library without CGO
RUN CGO_ENABLED=0 go install -a std

ENV APP_DIR $GOPATH{{.Appdir}}
RUN mkdir -p $APP_DIR

# Set the entrypoint
ENTRYPOINT $APP_DIR/{{.Entrypoint}}
ADD . $APP_DIR

# Compile the binary and statically link
RUN cd $APP_DIR
RUN CGO_ENABLED=0 godep go build -ldflags '-d -w -s'

EXPOSE {{.Expose}}
`

// Dockerfile holds the information about the Docker container.
type Dockerfile struct {
	BaseImage  string
	Appdir     string
	Entrypoint string
	Expose     string
}

var (
	expose    string
	baseImage string
)

func init() {
	fs := flag.NewFlagSet("dockerize", flag.ContinueOnError)
	fs.StringVar(&baseImage, "image", "", "Sets the base image of the Docker container.")
	fs.StringVar(&expose, "expose", "8080", "Port(s) to expose in the Docker container.")
	cmdDockerize.Flag = *fs
}

func dockerizeApp(cmd *Command, args []string) int {
	if err := cmd.Flag.Parse(args); err != nil {
		logger.Fatalf("Error parsing flags: %v", err.Error())
	}

	logger.Info("Generating Dockerfile...")

	gopath := os.Getenv("GOPATH")
	dir, err := filepath.Abs(".")
	MustCheck(err)

	if len(baseImage) == 0 {
		baseImage = "library/golang:latest"
	}

	appdir := strings.Replace(dir, gopath, "", 1)

	// In case of multiple ports to expose inside the container,
	// replace all the commas with whitespaces.
	// See the verb EXPOSE in the Docker documentation.
	if strings.Contains(expose, ",") {
		expose = strings.Replace(expose, ",", " ", -1)
	}

	_, entrypoint := path.Split(appdir)
	dockerfile := Dockerfile{
		BaseImage:  baseImage,
		Appdir:     appdir,
		Entrypoint: entrypoint,
		Expose:     expose,
	}

	generateDockerfile(dockerfile)
	return 0
}

func generateDockerfile(df Dockerfile) {
	t := template.Must(template.New("dockerBuildTemplate").Parse(dockerBuildTemplate)).Funcs(BeeFuncMap())

	f, err := os.Create("Dockerfile")
	if err != nil {
		logger.Fatalf("Error writing Dockerfile: %v", err.Error())
	}
	defer CloseFile(f)

	t.Execute(f, df)

	logger.Success("Dockerfile generated.")
}
<|endoftext|>


<commit_before>package Dictionary

import "strings"
import "sort"

type KeyValuePair struct {
	Key       string
	Lower_key string
	Value     string
}

type Dictionary map[string]*KeyValuePair

func NewDictionary() Dictionary {
	return make(Dictionary)
}

func (this *Dictionary) Get(key string) string {
	lower_key := strings.ToLower(key)

	node, ok := (*this)[lower_key]
	if ok {
		return node.Value
	} else {
		return ""
	}
}

func (this *Dictionary) Remove(key string) {
	delete(*this,strings.ToLower(key))
}

func (this *Dictionary) Set(key string, value string) {
	if value == "" {
		this.Remove(key)
		return
	}
	tmp := new(KeyValuePair)
	tmp.Key = key
	tmp.Lower_key = strings.ToLower(key)
	tmp.Value = value
	(*this)[tmp.Lower_key] = tmp
}

func (this *Dictionary) Iter() chan *KeyValuePair {
	ch := make(chan *KeyValuePair, 0)
	go func(this *Dictionary) {
		for _, pair := range *this {
			ch <- pair
		}
		close(ch)
		return
	}(this)
	return ch
}

type KeyValueList struct {
	list []KeyValuePair
}

func (this *KeyValueList) Len() int {
	return len(this.list)
}

func (this *KeyValueList) Less(i, j int) bool {
	return this.list[i].Key < this.list[j].Key
}

func (this *KeyValueList) Swap(i, j int) {
	this.list[i], this.list[j] = this.list[j], this.list[i]
}

func (this *Dictionary) ToArray() []KeyValuePair {
	list := []KeyValuePair{}
	for _, pair := range *this {
		list = append(list, *pair)
	}
	return list
}

func (this *Dictionary) SortIter() chan KeyValuePair {
	array := new(KeyValueList)
	array.list = this.ToArray()
	sort.Sort(array)
	ch := make(chan KeyValuePair)
	go func() {
		for i := 0; i < len(array.list); i++ {
			ch <- array.list[i]
		}
		close(ch)
	}()
	return ch
}
<commit_msg>Change Dictionary's type from plain map to struct<commit_after>package Dictionary

import "strings"
import "sort"

type KeyValuePair struct {
	Key       string
	Lower_key string
	Value     string
}

type Dictionary struct {
	m map[string]*KeyValuePair
}

func NewDictionary() *Dictionary {
	this := new(Dictionary)
	this.m = make(map[string]*KeyValuePair)
	return this
}

func (this *Dictionary) Get(key string) string {
	lower_key := strings.ToLower(key)

	node, ok := this.m[lower_key]
	if ok {
		return node.Value
	} else {
		return ""
	}
}

func (this *Dictionary) Remove(key string) {
	delete(this.m,strings.ToLower(key))
}

func (this *Dictionary) Set(key string, value string) {
	if value == "" {
		this.Remove(key)
		return
	}
	tmp := new(KeyValuePair)
	tmp.Key = key
	tmp.Lower_key = strings.ToLower(key)
	tmp.Value = value
	this.m[tmp.Lower_key] = tmp
}

func (this *Dictionary) Iter() chan *KeyValuePair {
	ch := make(chan *KeyValuePair, 0)
	go func() {
		for _, pair := range this.m {
			ch <- pair
		}
		close(ch)
		return
	}()
	return ch
}

type KeyValueList struct {
	list []KeyValuePair
}

func (this *KeyValueList) Len() int {
	return len(this.list)
}

func (this *KeyValueList) Less(i, j int) bool {
	return this.list[i].Key < this.list[j].Key
}

func (this *KeyValueList) Swap(i, j int) {
	this.list[i], this.list[j] = this.list[j], this.list[i]
}

func (this *Dictionary) ToArray() []KeyValuePair {
	list := []KeyValuePair{}
	for _, pair := range this.m {
		list = append(list, *pair)
	}
	return list
}

func (this *Dictionary) SortIter() chan KeyValuePair {
	array := new(KeyValueList)
	array.list = this.ToArray()
	sort.Sort(array)
	ch := make(chan KeyValuePair)
	go func() {
		for i := 0; i < len(array.list); i++ {
			ch <- array.list[i]
		}
		close(ch)
	}()
	return ch
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/Xe/x/internal"
	"github.com/Xe/x/markov"
	"github.com/Xe/x/web/switchcounter"
	"github.com/bwmarrin/discordgo"
	"github.com/joeshaw/envdecode"
	_ "github.com/joho/godotenv/autoload"
	"github.com/peterh/liner"
)

var (
	repl = flag.Bool("repl", false, "open a bot repl in the console?")
)

// lipuSona is the configuration.
type lipuSona struct {
	DiscordToken            string   `env:"DISCORD_TOKEN,required"` // lipu pi lukin ala
	TokiPonaTokenizerAPIURL string   `env:"TOKI_PONA_TOKENIZER_API_URL,default=https://us-central1-golden-cove-408.cloudfunctions.net/function-1"`
	SwitchCounterWebhook    string   `env:"SWITCH_COUNTER_WEBHOOK,required"`
	IloNimi                 string   `env:"ILO_NIMI,default=Kesi"`
	JanLawa                 []string `env:"JAN_LAWA,required"`
}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	flag.Parse()
	internal.HandleLicense()

	var cfg lipuSona
	err := envdecode.StrictDecode(&cfg)
	if err != nil {
		log.Fatal(err)
	}
	cfg.JanLawa = append(cfg.JanLawa, "console")

	flag.Parse()
	internal.HandleLicense()

	//pretty.Println(cfg)

	sw := switchcounter.NewHTTPClient(http.DefaultClient, cfg.SwitchCounterWebhook)

	line := liner.NewLiner()
	defer line.Close()

	chain := markov.NewChain(3)
	err = chain.Load("cadey.gob")
	if err != nil {
		log.Fatal(err)
	}

	words, err := loadWords("./tokipona.json")
	if err != nil {
		log.Fatal(err)
	}

	i := ilo{
		cfg:   cfg,
		sw:    sw,
		chain: chain,
		words: words,
	}

	line.SetCtrlCAborts(true)

	omc := func(s *discordgo.Session, m *discordgo.MessageCreate) {
		// Ignore all messages created by the bot itself
		// This isn't required in this specific example but it's a good practice.
		if m.Author.ID == s.State.User.ID {
			return
		}

		mentionsMe := false
		for _, us := range m.Mentions {
			if us.ID == s.State.User.ID {
				mentionsMe = true
				break
			}
		}

		if !mentionsMe {
			return
		}

		s.ChannelMessageSend(m.ChannelID, chain.Generate(15))
	}

	mc := func(s *discordgo.Session, m *discordgo.MessageCreate) {
		// Ignore all messages created by the bot itself
		// This isn't required in this specific example but it's a good practice.
		if m.Author.ID == s.State.User.ID {
			return
		}

		msg := m.ContentWithMentionsReplaced()
		if !i.tokiNiTokiPonaAnuSeme(msg) {
			return
		}

		result, err := i.parse(m.Author.ID, msg)
		if err != nil {
			switch err {
			case ErrJanLawaAla, ErrUnknownAction:
				s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("mi ken ala la %v", err))
				return
			}

			log.Printf("other error: %s", err)
			return
		}

		s.ChannelMessageSend(m.ChannelID, result.msg)
	}

	if *repl {
		for {
			if inp, err := line.Prompt("|lipu: "); err == nil {
				if inp == "" {
					return
				}

				line.AppendHistory(inp)

				result, err := i.parse("console", inp)
				if err != nil {
					log.Printf("error: %v", err)
					continue
				}

				fmt.Println(result.msg)
			} else if err == liner.ErrPromptAborted {
				log.Print("Aborted")
				break
			} else {
				log.Print("Error reading line: ", err)
				break
			}
		}

		os.Exit(0)
	} else {
		dg, err := discordgo.New("Bot " + cfg.DiscordToken)
		if err != nil {
			log.Fatal(err)
		}

		dg.AddHandler(mc)
		dg.AddHandler(omc)
		err = dg.Open()
		if err != nil {
			log.Fatal(err)
		}
		defer dg.Close()

		log.Println("bot is running")
		sc := make(chan os.Signal, 1)
		signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
		<-sc
	}
}
<commit_msg>discord/ilo-kesi: mention people less<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"syscall"
	"time"

	"github.com/Xe/x/internal"
	"github.com/Xe/x/markov"
	"github.com/Xe/x/web/switchcounter"
	"github.com/bwmarrin/discordgo"
	"github.com/joeshaw/envdecode"
	_ "github.com/joho/godotenv/autoload"
	"github.com/peterh/liner"
)

var (
	repl = flag.Bool("repl", false, "open a bot repl in the console?")
)

// lipuSona is the configuration.
type lipuSona struct {
	DiscordToken            string   `env:"DISCORD_TOKEN,required"` // lipu pi lukin ala
	TokiPonaTokenizerAPIURL string   `env:"TOKI_PONA_TOKENIZER_API_URL,default=https://us-central1-golden-cove-408.cloudfunctions.net/function-1"`
	SwitchCounterWebhook    string   `env:"SWITCH_COUNTER_WEBHOOK,required"`
	IloNimi                 string   `env:"ILO_NIMI,default=Kesi"`
	JanLawa                 []string `env:"JAN_LAWA,required"`
}

func init() {
	rand.Seed(time.Now().UnixNano())
}

var mentionRex = regexp.MustCompile(`<((@!?\d+)|(:.+?:\d+))>`)

func main() {
	flag.Parse()
	internal.HandleLicense()

	var cfg lipuSona
	err := envdecode.StrictDecode(&cfg)
	if err != nil {
		log.Fatal(err)
	}
	cfg.JanLawa = append(cfg.JanLawa, "console")

	flag.Parse()
	internal.HandleLicense()

	//pretty.Println(cfg)

	sw := switchcounter.NewHTTPClient(http.DefaultClient, cfg.SwitchCounterWebhook)

	line := liner.NewLiner()
	defer line.Close()

	chain := markov.NewChain(3)
	err = chain.Load("cadey.gob")
	if err != nil {
		log.Fatal(err)
	}

	words, err := loadWords("./tokipona.json")
	if err != nil {
		log.Fatal(err)
	}

	i := ilo{
		cfg:   cfg,
		sw:    sw,
		chain: chain,
		words: words,
	}

	line.SetCtrlCAborts(true)

	omc := func(s *discordgo.Session, m *discordgo.MessageCreate) {
		// Ignore all messages created by the bot itself
		// This isn't required in this specific example but it's a good practice.
		if m.Author.ID == s.State.User.ID {
			return
		}

		mentionsMe := false
		for _, us := range m.Mentions {
			if us.ID == s.State.User.ID {
				mentionsMe = true
				break
			}
		}

		if !mentionsMe {
			return
		}

		s.ChannelMessageSend(m.ChannelID, mentionRex.ReplaceAllString(chain.Generate(15), "<narvi'a>"))
	}

	mc := func(s *discordgo.Session, m *discordgo.MessageCreate) {
		// Ignore all messages created by the bot itself
		// This isn't required in this specific example but it's a good practice.
		if m.Author.ID == s.State.User.ID {
			return
		}

		msg := m.ContentWithMentionsReplaced()
		if !i.tokiNiTokiPonaAnuSeme(msg) {
			return
		}

		result, err := i.parse(m.Author.ID, msg)
		if err != nil {
			switch err {
			case ErrJanLawaAla, ErrUnknownAction:
				s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("mi ken ala la %v", err))
				return
			}

			log.Printf("other error: %s", err)
			return
		}

		s.ChannelMessageSend(m.ChannelID, result.msg)
	}

	if *repl {
		for {
			if inp, err := line.Prompt("|lipu: "); err == nil {
				if inp == "" {
					return
				}

				line.AppendHistory(inp)

				result, err := i.parse("console", inp)
				if err != nil {
					log.Printf("error: %v", err)
					continue
				}

				fmt.Println(result.msg)
			} else if err == liner.ErrPromptAborted {
				log.Print("Aborted")
				break
			} else {
				log.Print("Error reading line: ", err)
				break
			}
		}

		os.Exit(0)
	} else {
		dg, err := discordgo.New("Bot " + cfg.DiscordToken)
		if err != nil {
			log.Fatal(err)
		}

		dg.AddHandler(mc)
		dg.AddHandler(omc)
		err = dg.Open()
		if err != nil {
			log.Fatal(err)
		}
		defer dg.Close()

		log.Println("bot is running")
		sc := make(chan os.Signal, 1)
		signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
		<-sc
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"errors"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"
)

// ClusterConfig is a struct to configure the default cluster implementation
// of gocoql. It has a varity of attributes that can be used to modify the
// behavior to fit the most common use cases. Applications that requre a
// different setup must implement their own cluster.
type ClusterConfig struct {
	Hosts        []string      // addresses for the initial connections
	CQLVersion   string        // CQL version (default: 3.0.0)
	ProtoVersion int           // version of the native protocol (default: 2)
	Timeout      time.Duration // connection timeout (default: 200ms)
	DefaultPort  int           // default port (default: 9042)
	Keyspace     string        // initial keyspace (optional)
	NumConns     int           // number of connections per host (default: 2)
	NumStreams   int           // number of streams per connection (default: 128)
	DelayMin     time.Duration // minimum reconnection delay (default: 1s)
	DelayMax     time.Duration // maximum reconnection delay (default: 10min)
	StartupMin   int           // wait for StartupMin hosts (default: len(Hosts)/2+1)
	Consistency  Consistency   // default consistency level (default: Quorum)
	Compressor   Compressor    // compression algorithm (default: nil)
}

// NewCluster generates a new config for the default cluster implementation.
func NewCluster(hosts ...string) *ClusterConfig {
	cfg := &ClusterConfig{
		Hosts:        hosts,
		CQLVersion:   "3.0.0",
		ProtoVersion: 2,
		Timeout:      200 * time.Millisecond,
		DefaultPort:  9042,
		NumConns:     2,
		NumStreams:   128,
		DelayMin:     1 * time.Second,
		DelayMax:     10 * time.Minute,
		StartupMin:   len(hosts)/2 + 1,
		Consistency:  Quorum,
	}
	return cfg
}

// CreateSession initializes the cluster based on this config and returns a
// session object that can be used to interact with the database.
func (cfg *ClusterConfig) CreateSession() (*Session, error) {

	//Check that hosts in the ClusterConfig is not empty
	if len(cfg.Hosts) < 1 {
		return nil, ErrNoHosts
	}

	impl := &clusterImpl{
		cfg:      *cfg,
		hostPool: NewRoundRobin(),
		connPool: make(map[string]*RoundRobin),
		conns:    make(map[*Conn]struct{}),
		quitWait: make(chan bool),
		keyspace: cfg.Keyspace,
	}
	impl.wgStart.Add(1)
	for i := 0; i < len(impl.cfg.Hosts); i++ {
		addr := strings.TrimSpace(impl.cfg.Hosts[i])
		if strings.Index(addr, ":") < 0 {
			addr = fmt.Sprintf("%s:%d", addr, impl.cfg.DefaultPort)
		}
		for j := 0; j < impl.cfg.NumConns; j++ {
			go impl.connect(addr)
		}
	}
	impl.wgStart.Wait()
	s := NewSession(impl)
	s.SetConsistency(cfg.Consistency)
	return s, nil
}

type clusterImpl struct {
	cfg      ClusterConfig
	hostPool *RoundRobin
	connPool map[string]*RoundRobin
	conns    map[*Conn]struct{}
	keyspace string
	mu       sync.Mutex

	started bool
	wgStart sync.WaitGroup

	quit     bool
	quitWait chan bool
	quitOnce sync.Once
}

func (c *clusterImpl) connect(addr string) {
	cfg := ConnConfig{
		ProtoVersion: c.cfg.ProtoVersion,
		CQLVersion:   c.cfg.CQLVersion,
		Timeout:      c.cfg.Timeout,
		NumStreams:   c.cfg.NumStreams,
		Compressor:   c.cfg.Compressor,
	}
	delay := c.cfg.DelayMin
	for {
		conn, err := Connect(addr, cfg, c)
		if err != nil {
			log.Printf("failed to connect to %q: %v", addr, err)
			select {
			case <-time.After(delay):
				if delay *= 2; delay > c.cfg.DelayMax {
					delay = c.cfg.DelayMax
				}
				continue
			case <-c.quitWait:
				return
			}
		}
		c.addConn(conn, "")
		return
	}
}

func (c *clusterImpl) changeKeyspace(conn *Conn, keyspace string, connected bool) {
	if err := conn.UseKeyspace(keyspace); err != nil {
		conn.Close()
		if connected {
			c.removeConn(conn)
		}
		go c.connect(conn.Address())
	}
	if !connected {
		c.addConn(conn, keyspace)
	}
}

func (c *clusterImpl) addConn(conn *Conn, keyspace string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.quit {
		conn.Close()
		return
	}
	if keyspace != c.keyspace && c.keyspace != "" {
		// change the keyspace before adding the node to the pool
		go c.changeKeyspace(conn, c.keyspace, false)
		return
	}
	connPool := c.connPool[conn.Address()]
	if connPool == nil {
		connPool = NewRoundRobin()
		c.connPool[conn.Address()] = connPool
		c.hostPool.AddNode(connPool)
		if !c.started && c.hostPool.Size() >= c.cfg.StartupMin {
			c.started = true
			c.wgStart.Done()
		}
	}
	connPool.AddNode(conn)
	c.conns[conn] = struct{}{}
}

func (c *clusterImpl) removeConn(conn *Conn) {
	c.mu.Lock()
	defer c.mu.Unlock()
	conn.Close()
	connPool := c.connPool[conn.addr]
	if connPool == nil {
		return
	}
	connPool.RemoveNode(conn)
	if connPool.Size() == 0 {
		c.hostPool.RemoveNode(connPool)
	}
	delete(c.conns, conn)
}

func (c *clusterImpl) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		// ignore all non-fatal errors
		return
	}
	c.removeConn(conn)
	if !c.quit {
		go c.connect(conn.Address()) // reconnect
	}
}

func (c *clusterImpl) HandleKeyspace(conn *Conn, keyspace string) {
	c.mu.Lock()
	if c.keyspace == keyspace {
		c.mu.Unlock()
		return
	}
	c.keyspace = keyspace
	conns := make([]*Conn, 0, len(c.conns))
	for conn := range c.conns {
		conns = append(conns, conn)
	}
	c.mu.Unlock()

	// change the keyspace of all other connections too
	for i := 0; i < len(conns); i++ {
		if conns[i] == conn {
			continue
		}
		c.changeKeyspace(conns[i], keyspace, true)
	}
}

func (c *clusterImpl) Pick(qry *Query) *Conn {
	return c.hostPool.Pick(qry)
}

func (c *clusterImpl) Close() {
	c.quitOnce.Do(func() {
		c.mu.Lock()
		defer c.mu.Unlock()
		c.quit = true
		close(c.quitWait)
		for conn := range c.conns {
			conn.Close()
		}
	})
}

var (
	ErrNoHosts = errors.New("no hosts provided")
)
<commit_msg>Fix #40 pool/reconnection issue as suggested by @Kimbsen<commit_after>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"errors"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"
)

// ClusterConfig is a struct to configure the default cluster implementation
// of gocoql. It has a varity of attributes that can be used to modify the
// behavior to fit the most common use cases. Applications that requre a
// different setup must implement their own cluster.
type ClusterConfig struct {
	Hosts        []string      // addresses for the initial connections
	CQLVersion   string        // CQL version (default: 3.0.0)
	ProtoVersion int           // version of the native protocol (default: 2)
	Timeout      time.Duration // connection timeout (default: 200ms)
	DefaultPort  int           // default port (default: 9042)
	Keyspace     string        // initial keyspace (optional)
	NumConns     int           // number of connections per host (default: 2)
	NumStreams   int           // number of streams per connection (default: 128)
	DelayMin     time.Duration // minimum reconnection delay (default: 1s)
	DelayMax     time.Duration // maximum reconnection delay (default: 10min)
	StartupMin   int           // wait for StartupMin hosts (default: len(Hosts)/2+1)
	Consistency  Consistency   // default consistency level (default: Quorum)
	Compressor   Compressor    // compression algorithm (default: nil)
}

// NewCluster generates a new config for the default cluster implementation.
func NewCluster(hosts ...string) *ClusterConfig {
	cfg := &ClusterConfig{
		Hosts:        hosts,
		CQLVersion:   "3.0.0",
		ProtoVersion: 2,
		Timeout:      200 * time.Millisecond,
		DefaultPort:  9042,
		NumConns:     2,
		NumStreams:   128,
		DelayMin:     1 * time.Second,
		DelayMax:     10 * time.Minute,
		StartupMin:   len(hosts)/2 + 1,
		Consistency:  Quorum,
	}
	return cfg
}

// CreateSession initializes the cluster based on this config and returns a
// session object that can be used to interact with the database.
func (cfg *ClusterConfig) CreateSession() (*Session, error) {

	//Check that hosts in the ClusterConfig is not empty
	if len(cfg.Hosts) < 1 {
		return nil, ErrNoHosts
	}

	impl := &clusterImpl{
		cfg:      *cfg,
		hostPool: NewRoundRobin(),
		connPool: make(map[string]*RoundRobin),
		conns:    make(map[*Conn]struct{}),
		quitWait: make(chan bool),
		keyspace: cfg.Keyspace,
	}
	impl.wgStart.Add(1)
	for i := 0; i < len(impl.cfg.Hosts); i++ {
		addr := strings.TrimSpace(impl.cfg.Hosts[i])
		if strings.Index(addr, ":") < 0 {
			addr = fmt.Sprintf("%s:%d", addr, impl.cfg.DefaultPort)
		}
		for j := 0; j < impl.cfg.NumConns; j++ {
			go impl.connect(addr)
		}
	}
	impl.wgStart.Wait()
	s := NewSession(impl)
	s.SetConsistency(cfg.Consistency)
	return s, nil
}

type clusterImpl struct {
	cfg      ClusterConfig
	hostPool *RoundRobin
	connPool map[string]*RoundRobin
	conns    map[*Conn]struct{}
	keyspace string
	mu       sync.Mutex

	started bool
	wgStart sync.WaitGroup

	quit     bool
	quitWait chan bool
	quitOnce sync.Once
}

func (c *clusterImpl) connect(addr string) {
	cfg := ConnConfig{
		ProtoVersion: c.cfg.ProtoVersion,
		CQLVersion:   c.cfg.CQLVersion,
		Timeout:      c.cfg.Timeout,
		NumStreams:   c.cfg.NumStreams,
		Compressor:   c.cfg.Compressor,
	}
	delay := c.cfg.DelayMin
	for {
		conn, err := Connect(addr, cfg, c)
		if err != nil {
			log.Printf("failed to connect to %q: %v", addr, err)
			select {
			case <-time.After(delay):
				if delay *= 2; delay > c.cfg.DelayMax {
					delay = c.cfg.DelayMax
				}
				continue
			case <-c.quitWait:
				return
			}
		}
		c.addConn(conn, "")
		return
	}
}

func (c *clusterImpl) changeKeyspace(conn *Conn, keyspace string, connected bool) {
	if err := conn.UseKeyspace(keyspace); err != nil {
		conn.Close()
		if connected {
			c.removeConn(conn)
		}
		go c.connect(conn.Address())
	}
	if !connected {
		c.addConn(conn, keyspace)
	}
}

func (c *clusterImpl) addConn(conn *Conn, keyspace string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.quit {
		conn.Close()
		return
	}
	if keyspace != c.keyspace && c.keyspace != "" {
		// change the keyspace before adding the node to the pool
		go c.changeKeyspace(conn, c.keyspace, false)
		return
	}
	connPool := c.connPool[conn.Address()]
	if connPool == nil {
		connPool = NewRoundRobin()
		c.connPool[conn.Address()] = connPool
		c.hostPool.AddNode(connPool)
		if !c.started && c.hostPool.Size() >= c.cfg.StartupMin {
			c.started = true
			c.wgStart.Done()
		}
	}
	connPool.AddNode(conn)
	c.conns[conn] = struct{}{}
}

func (c *clusterImpl) removeConn(conn *Conn) {
	c.mu.Lock()
	defer c.mu.Unlock()
	conn.Close()
	connPool := c.connPool[conn.addr]
	if connPool == nil {
		return
	}
	connPool.RemoveNode(conn)
	if connPool.Size() == 0 {
		c.hostPool.RemoveNode(connPool)
		delete(c.connPool, conn.addr)
	}
	delete(c.conns, conn)
}

func (c *clusterImpl) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		// ignore all non-fatal errors
		return
	}
	c.removeConn(conn)
	if !c.quit {
		go c.connect(conn.Address()) // reconnect
	}
}

func (c *clusterImpl) HandleKeyspace(conn *Conn, keyspace string) {
	c.mu.Lock()
	if c.keyspace == keyspace {
		c.mu.Unlock()
		return
	}
	c.keyspace = keyspace
	conns := make([]*Conn, 0, len(c.conns))
	for conn := range c.conns {
		conns = append(conns, conn)
	}
	c.mu.Unlock()

	// change the keyspace of all other connections too
	for i := 0; i < len(conns); i++ {
		if conns[i] == conn {
			continue
		}
		c.changeKeyspace(conns[i], keyspace, true)
	}
}

func (c *clusterImpl) Pick(qry *Query) *Conn {
	return c.hostPool.Pick(qry)
}

func (c *clusterImpl) Close() {
	c.quitOnce.Do(func() {
		c.mu.Lock()
		defer c.mu.Unlock()
		c.quit = true
		close(c.quitWait)
		for conn := range c.conns {
			conn.Close()
		}
	})
}

var (
	ErrNoHosts = errors.New("no hosts provided")
)
<|endoftext|>