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

// CMAC message authentication code, defined in
// NIST Special Publication SP 800-38B.

package cmac

import (
	"crypto/cipher"
	"crypto/subtle"
	"errors"
	"hash"
)

const (
	// minimal irreducible polynomial of degree b
	r64  = 0x1b
	r128 = 0x87
)

type cmac struct {
	k1, k2, ci, digest []byte
	p                  int // position in ci
	c                  cipher.Block
}

// New returns a new instance of a CMAC message authentication code
// digest using the given cipher.Block.
func New(c cipher.Block) (hash.Hash, error) {
	var r int
	n := c.BlockSize()
	switch n {
	case 64 / 8:
		r = r64
	case 128 / 8:
		r = r128
	default:
		return nil, errors.New("NewCMAC: invalid cipher block size")
	}

	d := new(cmac)
	d.c = c
	d.k1 = make([]byte, n)
	d.k2 = make([]byte, n)
	d.ci = make([]byte, n)
	d.digest = make([]byte, n)

	// Subkey generation, p. 7
	c.Encrypt(d.k1, d.k1)

	v1 := shift1(d.k1, d.k1)
	d.k1[n-1] ^= byte(subtle.ConstantTimeSelect(v1, r, 0))

	v2 := shift1(d.k1, d.k2)
	d.k2[n-1] ^= byte(subtle.ConstantTimeSelect(v2, r, 0))

	return d, nil
}

// Reset clears the digest state, starting a new digest.
func (d *cmac) Reset() {
	for i := range d.ci {
		d.ci[i] = 0
	}
	d.p = 0
}

// Write adds the given data to the digest state.
func (d *cmac) Write(p []byte) (nn int, err error) {
	nn = len(p)
	bs := len(d.ci)
	left := bs - d.p

	if len(p) > left {
		xor(d.ci[d.p:], p[:left])
		p = p[left:]
		d.c.Encrypt(d.ci, d.ci)
		d.p = 0
	}

	for len(p) > bs {
		xor(d.ci, p[:bs])
		p = p[bs:]
		d.c.Encrypt(d.ci, d.ci)
	}

	if len(p) > 0 {
		xor(d.ci[d.p:], p)
		d.p += len(p)
	}
	return
}

// Sum returns the CMAC digest, one cipher block in length,
// of the data written with Write.
func (d *cmac) Sum(in []byte) []byte {
	// Finish last block, mix in key, encrypt.
	// Don't edit ci, in case caller wants
	// to keep digesting after call to Sum.
	k := d.k1
	if d.p < len(d.digest) {
		k = d.k2
	}
	for i := 0; i < len(d.ci); i++ {
		d.digest[i] = d.ci[i] ^ k[i]
	}
	if d.p < len(d.digest) {
		d.digest[d.p] ^= 0x80
	}
	d.c.Encrypt(d.digest, d.digest)
	return append(in, d.digest...)
}

func (d *cmac) Size() int { return len(d.digest) }

func (d *cmac) BlockSize() int { return d.c.BlockSize() }

func shift1(src, dst []byte) int {
	var b byte
	for i := len(src) - 1; i >= 0; i-- {
		bb := src[i] >> 7
		dst[i] = src[i]<<1 | b
		b = bb
	}
	return int(b)
}

func xor(a, b []byte) {
	for i, v := range b {
		a[i] ^= v
	}
}
<commit_msg>Fix error message<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// CMAC message authentication code, defined in
// NIST Special Publication SP 800-38B.

package cmac

import (
	"crypto/cipher"
	"crypto/subtle"
	"errors"
	"hash"
)

const (
	// minimal irreducible polynomial of degree b
	r64  = 0x1b
	r128 = 0x87
)

type cmac struct {
	k1, k2, ci, digest []byte
	p                  int // position in ci
	c                  cipher.Block
}

// New returns a new instance of a CMAC message authentication code
// digest using the given cipher.Block.
func New(c cipher.Block) (hash.Hash, error) {
	var r int
	n := c.BlockSize()
	switch n {
	case 64 / 8:
		r = r64
	case 128 / 8:
		r = r128
	default:
		return nil, errors.New("cmac: invalid cipher block size")
	}

	d := new(cmac)
	d.c = c
	d.k1 = make([]byte, n)
	d.k2 = make([]byte, n)
	d.ci = make([]byte, n)
	d.digest = make([]byte, n)

	// Subkey generation, p. 7
	c.Encrypt(d.k1, d.k1)

	v1 := shift1(d.k1, d.k1)
	d.k1[n-1] ^= byte(subtle.ConstantTimeSelect(v1, r, 0))

	v2 := shift1(d.k1, d.k2)
	d.k2[n-1] ^= byte(subtle.ConstantTimeSelect(v2, r, 0))

	return d, nil
}

// Reset clears the digest state, starting a new digest.
func (d *cmac) Reset() {
	for i := range d.ci {
		d.ci[i] = 0
	}
	d.p = 0
}

// Write adds the given data to the digest state.
func (d *cmac) Write(p []byte) (nn int, err error) {
	nn = len(p)
	bs := len(d.ci)
	left := bs - d.p

	if len(p) > left {
		xor(d.ci[d.p:], p[:left])
		p = p[left:]
		d.c.Encrypt(d.ci, d.ci)
		d.p = 0
	}

	for len(p) > bs {
		xor(d.ci, p[:bs])
		p = p[bs:]
		d.c.Encrypt(d.ci, d.ci)
	}

	if len(p) > 0 {
		xor(d.ci[d.p:], p)
		d.p += len(p)
	}
	return
}

// Sum returns the CMAC digest, one cipher block in length,
// of the data written with Write.
func (d *cmac) Sum(in []byte) []byte {
	// Finish last block, mix in key, encrypt.
	// Don't edit ci, in case caller wants
	// to keep digesting after call to Sum.
	k := d.k1
	if d.p < len(d.digest) {
		k = d.k2
	}
	for i := 0; i < len(d.ci); i++ {
		d.digest[i] = d.ci[i] ^ k[i]
	}
	if d.p < len(d.digest) {
		d.digest[d.p] ^= 0x80
	}
	d.c.Encrypt(d.digest, d.digest)
	return append(in, d.digest...)
}

func (d *cmac) Size() int { return len(d.digest) }

func (d *cmac) BlockSize() int { return d.c.BlockSize() }

func shift1(src, dst []byte) int {
	var b byte
	for i := len(src) - 1; i >= 0; i-- {
		bb := src[i] >> 7
		dst[i] = src[i]<<1 | b
		b = bb
	}
	return int(b)
}

func xor(a, b []byte) {
	for i, v := range b {
		a[i] ^= v
	}
}
<|endoftext|>


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

// +build 386 arm arm64 mips64 mips64le ppc64 ppc64le

package crc32

// The file contains the generic version of updateCastagnoli which just calls
// the software implementation.

func updateCastagnoli(crc uint32, p []byte) uint32 {
	return update(crc, castagnoliTable, p)
}

func updateIEEE(crc uint32, p []byte) uint32 {
	// only use slicing-by-8 when input is >= 4KB
	if len(p) >= 4096 {
		ieeeTable8Once.Do(func() {
			ieeeTable8 = makeTable8(IEEE)
		})
		return updateSlicingBy8(crc, ieeeTable8, p)
	}
	return update(crc, IEEETable, p)
}
<commit_msg>hash/crc32: add sparc64 build tags<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build 386 arm arm64 mips64 mips64le ppc64 ppc64le sparc64

package crc32

// The file contains the generic version of updateCastagnoli which just calls
// the software implementation.

func updateCastagnoli(crc uint32, p []byte) uint32 {
	return update(crc, castagnoliTable, p)
}

func updateIEEE(crc uint32, p []byte) uint32 {
	// only use slicing-by-8 when input is >= 4KB
	if len(p) >= 4096 {
		ieeeTable8Once.Do(func() {
			ieeeTable8 = makeTable8(IEEE)
		})
		return updateSlicingBy8(crc, ieeeTable8, p)
	}
	return update(crc, IEEETable, p)
}
<|endoftext|>


<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package util

import (
	"github.com/getgauge/gauge/config"
	. "gopkg.in/check.v1"
	"path/filepath"
)

func (s *MySuite) TestPrefixingMessage(c *C) {
	prefixedLines := AddPrefixToEachLine("Hello\nWorld", "[my-plugin Plugin] : ")
	c.Assert(prefixedLines, Equals, "[my-plugin Plugin] : Hello\n"+
		"[my-plugin Plugin] : World")
}

func (s *MySuite) TestPrefixingMessageEndingWithNewLine(c *C) {
	prefixedLines := AddPrefixToEachLine("Hello\nWorld\n", "[my-plugin Plugin] : ")
	c.Assert(prefixedLines, Equals, "[my-plugin Plugin] : Hello\n"+
		"[my-plugin Plugin] : World\n")

}

func (s *MySuite) TestPrefixingMultiLineMessagWithNewLine(c *C) {
	prefixedLines := AddPrefixToEachLine("\nHello\nWorld\n\nFoo bar\n", "[my-plugin Plugin] : ")
	c.Assert(prefixedLines, Equals, "[my-plugin Plugin] : \n"+
		"[my-plugin Plugin] : Hello\n"+
		"[my-plugin Plugin] : World\n"+
		"[my-plugin Plugin] : \n"+
		"[my-plugin Plugin] : Foo bar\n")

}

func (s *MySuite) TestGetPathToFile(c *C) {
	var path string
	config.ProjectRoot = "PROJECT_ROOT"

	path = GetPathToFile(string(filepath.Separator) + "resources")
	c.Assert(path, Equals, string(filepath.Separator) + "resources")

	path = GetPathToFile("resources")
	c.Assert(path, Equals, filepath.Join(config.ProjectRoot,"resources"))
}
<commit_msg>fixing util tests<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package util

import (
	"github.com/getgauge/gauge/config"
	. "gopkg.in/check.v1"
	"path/filepath"
)

func (s *MySuite) TestPrefixingMessage(c *C) {
	prefixedLines := AddPrefixToEachLine("Hello\nWorld", "[my-plugin Plugin] : ")
	c.Assert(prefixedLines, Equals, "[my-plugin Plugin] : Hello\n"+
		"[my-plugin Plugin] : World")
}

func (s *MySuite) TestPrefixingMessageEndingWithNewLine(c *C) {
	prefixedLines := AddPrefixToEachLine("Hello\nWorld\n", "[my-plugin Plugin] : ")
	c.Assert(prefixedLines, Equals, "[my-plugin Plugin] : Hello\n"+
		"[my-plugin Plugin] : World\n")

}

func (s *MySuite) TestPrefixingMultiLineMessagWithNewLine(c *C) {
	prefixedLines := AddPrefixToEachLine("\nHello\nWorld\n\nFoo bar\n", "[my-plugin Plugin] : ")
	c.Assert(prefixedLines, Equals, "[my-plugin Plugin] : \n"+
		"[my-plugin Plugin] : Hello\n"+
		"[my-plugin Plugin] : World\n"+
		"[my-plugin Plugin] : \n"+
		"[my-plugin Plugin] : Foo bar\n")

}

func (s *MySuite) TestGetPathToFile(c *C) {
	var path string
	config.ProjectRoot = "PROJECT_ROOT"
	absPath, _ := filepath.Abs("resources")
	path = GetPathToFile(absPath)
	c.Assert(path, Equals, absPath)

	path = GetPathToFile("resources")
	c.Assert(path, Equals, filepath.Join(config.ProjectRoot,"resources"))
}
<|endoftext|>


<commit_before>package util

import (
	"bytes"
	"math/rand"
	"testing"

	mh "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multihash"
)

func TestKey(t *testing.T) {

	h1, err := mh.Sum([]byte("beep boop"), mh.SHA2_256, -1)
	if err != nil {
		t.Error(err)
	}

	k1 := Key(h1)
	h2 := mh.Multihash(k1)
	k2 := Key(h2)

	if !bytes.Equal(h1, h2) {
		t.Error("Multihashes not equal.")
	}

	if k1 != k2 {
		t.Error("Keys not equal.")
	}
}

func TestByteChanReader(t *testing.T) {

	var data bytes.Buffer
	var data2 bytes.Buffer
	dch := make(chan []byte, 8)
	randr := NewTimeSeededRand()

	go func() {
		defer close(dch)
		for i := 0; i < rand.Intn(100)+100; i++ {
			chunk := make([]byte, rand.Intn(100000)+10)
			randr.Read(chunk)
			data.Write(chunk)
			// fmt.Printf("chunk: %6.d %v\n", len(chunk), chunk[:10])
			dch <- chunk
		}
	}()

	read := NewByteChanReader(dch)

	// read in random, weird sizes to exercise saving buffer.
	for {
		buf := make([]byte, rand.Intn(10)*10)
		n, err := read.Read(buf)
		data2.Write(buf[:n])
		// fmt.Printf("read: %6.d\n", n)
		if err != nil {
			break
		}
	}

	// fmt.Printf("lens: %d == %d\n", len(out), len(data.Bytes()))
	if !bytes.Equal(data2.Bytes(), data.Bytes()) {
		t.Fatal("Reader failed to stream correct bytes")
	}
}

func TestXOR(t *testing.T) {
	cases := [][3][]byte{
		[3][]byte{
			[]byte{0xFF, 0xFF, 0xFF},
			[]byte{0xFF, 0xFF, 0xFF},
			[]byte{0x00, 0x00, 0x00},
		},
		[3][]byte{
			[]byte{0x00, 0xFF, 0x00},
			[]byte{0xFF, 0xFF, 0xFF},
			[]byte{0xFF, 0x00, 0xFF},
		},
		[3][]byte{
			[]byte{0x55, 0x55, 0x55},
			[]byte{0x55, 0xFF, 0xAA},
			[]byte{0x00, 0xAA, 0xFF},
		},
	}

	for _, c := range cases {
		r := XOR(c[0], c[1])
		if !bytes.Equal(r, c[2]) {
			t.Error("XOR failed")
		}
	}
}
<commit_msg>remove comment<commit_after>package util

import (
	"bytes"
	"math/rand"
	"testing"

	mh "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multihash"
)

func TestKey(t *testing.T) {

	h1, err := mh.Sum([]byte("beep boop"), mh.SHA2_256, -1)
	if err != nil {
		t.Error(err)
	}

	k1 := Key(h1)
	h2 := mh.Multihash(k1)
	k2 := Key(h2)

	if !bytes.Equal(h1, h2) {
		t.Error("Multihashes not equal.")
	}

	if k1 != k2 {
		t.Error("Keys not equal.")
	}
}

func TestByteChanReader(t *testing.T) {

	var data bytes.Buffer
	var data2 bytes.Buffer
	dch := make(chan []byte, 8)
	randr := NewTimeSeededRand()

	go func() {
		defer close(dch)
		for i := 0; i < rand.Intn(100)+100; i++ {
			chunk := make([]byte, rand.Intn(100000)+10)
			randr.Read(chunk)
			data.Write(chunk)
			dch <- chunk
		}
	}()

	read := NewByteChanReader(dch)

	// read in random, weird sizes to exercise saving buffer.
	for {
		buf := make([]byte, rand.Intn(10)*10)
		n, err := read.Read(buf)
		data2.Write(buf[:n])
		if err != nil {
			break
		}
	}

	if !bytes.Equal(data2.Bytes(), data.Bytes()) {
		t.Fatal("Reader failed to stream correct bytes")
	}
}

func TestXOR(t *testing.T) {
	cases := [][3][]byte{
		[3][]byte{
			[]byte{0xFF, 0xFF, 0xFF},
			[]byte{0xFF, 0xFF, 0xFF},
			[]byte{0x00, 0x00, 0x00},
		},
		[3][]byte{
			[]byte{0x00, 0xFF, 0x00},
			[]byte{0xFF, 0xFF, 0xFF},
			[]byte{0xFF, 0x00, 0xFF},
		},
		[3][]byte{
			[]byte{0x55, 0x55, 0x55},
			[]byte{0x55, 0xFF, 0xAA},
			[]byte{0x00, 0xAA, 0xFF},
		},
	}

	for _, c := range cases {
		r := XOR(c[0], c[1])
		if !bytes.Equal(r, c[2]) {
			t.Error("XOR failed")
		}
	}
}
<|endoftext|>


<commit_before>package neptulon

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

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

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

// Conn is a client connection.
type Conn struct {
	ID           string     // Randomly generated unique client connection ID.
	Session      *cmap.CMap // Thread-safe data store for storing arbitrary data for this connection session.
	middleware   []func(ctx *ReqCtx) error
	resRoutes    *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws           *websocket.Conn
	wg           sync.WaitGroup
	deadline     time.Duration
	isClientConn bool
	connected    atomic.Value
}

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

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

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

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...Middleware) {
	for _, m := range middleware {
		c.MiddlewareFunc(m.Middleware)
	}
}

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

// Connect connects to the given WebSocket server.
// addr should be formatted as ws://host:port -or- wss://host:port (i.e. ws://127.0.0.1:3000 -or- wss://localhost:3000)
func (c *Conn) Connect(addr string) error {
	ws, err := websocket.Dial(addr, "", "http://localhost")
	if err != nil {
		return err
	}
	if err := c.setConn(ws); err != nil {
		return err
	}

	c.isClientConn = true
	c.wg.Add(1)
	go func() {
		defer recoverAndLog(c, &c.wg)
		c.startReceive()
	}()
	time.Sleep(time.Millisecond) // give receive goroutine a few cycles to start
	return nil
}

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

	return c.ws.RemoteAddr()
}

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

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

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

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

// Close closes the connection.
func (c *Conn) Close() error {
	c.connected.Store(false)
	if c.ws != nil {
		c.ws.Close()
	}
	return nil
}

// Wait waits for all message/connection handler goroutines to exit.
func (c *Conn) Wait() {
	c.wg.Wait()
}

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

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

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

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

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

// Reuse an established websocket.Conn.
func (c *Conn) setConn(ws *websocket.Conn) error {
	c.ws = ws
	c.connected.Store(true)
	if err := ws.SetDeadline(time.Now().Add(c.deadline)); err != nil {
		return fmt.Errorf("conn: error while setting websocket connection deadline: %v", err)
	}
	return nil
}

// startReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) startReceive() {
	defer c.Close()

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			// if we closed the connection
			if !c.connected.Load().(bool) {
				log.Printf("conn: closed %v: %v", c.ID, c.RemoteAddr())
				break
			}

			// if peer closed the connection
			if err == io.EOF {
				log.Printf("conn: peer disconnected %v: %v", c.ID, c.RemoteAddr())
				break
			}

			log.Printf("conn: error while receiving message: %v", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			c.wg.Add(1)
			go func() {
				defer recoverAndLog(c, &c.wg)
				if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
					log.Printf("conn: error while handling request: %v", err)
					c.Close()
				}
			}()

			continue
		}

		// if the message is not a JSON-RPC message
		if m.ID == "" || (m.Result == nil && m.Error == nil) {
			log.Printf("conn: received an unknown message %v: %v, %v", c.ID, c.RemoteAddr(), m)
			break
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			c.wg.Add(1)
			go func() {
				defer recoverAndLog(c, &c.wg)
				err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
				c.resRoutes.Delete(m.ID)
				if err != nil {
					log.Printf("conn: error while handling response: %v", err)
					c.Close()
				}
			}()
		} else {
			log.Printf("conn: error while handling response: got response to a request with unknown ID: %v", m.ID)
			break
		}
	}
}

func recoverAndLog(c *Conn, wg *sync.WaitGroup) {
	defer wg.Done()
	if err := recover(); err != nil {
		const size = 64 << 10
		buf := make([]byte, size)
		buf = buf[:runtime.Stack(buf, false)]
		log.Printf("conn: panic handling response %v: %v\n%s", c.RemoteAddr(), err, buf)
	}
}
<commit_msg>add Conn.DisconnHandler<commit_after>package neptulon

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

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

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

// Conn is a client connection.
type Conn struct {
	ID             string     // Randomly generated unique client connection ID.
	Session        *cmap.CMap // Thread-safe data store for storing arbitrary data for this connection session.
	middleware     []func(ctx *ReqCtx) error
	resRoutes      *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws             *websocket.Conn
	wg             sync.WaitGroup
	deadline       time.Duration
	isClientConn   bool
	connected      atomic.Value
	disconnHandler func(c *Conn)
}

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

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

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

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...Middleware) {
	for _, m := range middleware {
		c.MiddlewareFunc(m.Middleware)
	}
}

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

// DisconnHandler registers a function to handle disconnection event.
func (c *Conn) DisconnHandler(handler func(c *Conn)) {
	c.disconnHandler = handler
}

// Connect connects to the given WebSocket server.
// addr should be formatted as ws://host:port -or- wss://host:port (i.e. ws://127.0.0.1:3000 -or- wss://localhost:3000)
func (c *Conn) Connect(addr string) error {
	ws, err := websocket.Dial(addr, "", "http://localhost")
	if err != nil {
		return err
	}
	if err := c.setConn(ws); err != nil {
		return err
	}

	c.isClientConn = true
	c.wg.Add(1)
	go func() {
		defer recoverAndLog(c, &c.wg)
		c.startReceive()
	}()
	time.Sleep(time.Millisecond) // give receive goroutine a few cycles to start
	return nil
}

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

	return c.ws.RemoteAddr()
}

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

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

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

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

// Close closes the connection.
func (c *Conn) Close() error {
	c.connected.Store(false)
	if c.ws != nil {
		c.ws.Close()
	}
	return nil
}

// Wait waits for all message/connection handler goroutines to exit.
func (c *Conn) Wait() {
	c.wg.Wait()
}

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

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

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

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

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

// Reuse an established websocket.Conn.
func (c *Conn) setConn(ws *websocket.Conn) error {
	c.ws = ws
	c.connected.Store(true)
	if err := ws.SetDeadline(time.Now().Add(c.deadline)); err != nil {
		return fmt.Errorf("conn: error while setting websocket connection deadline: %v", err)
	}
	return nil
}

// startReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) startReceive() {
	defer c.disconnHandler(c)
	defer c.Close()

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			// if we closed the connection
			if !c.connected.Load().(bool) {
				log.Printf("conn: closed %v: %v", c.ID, c.RemoteAddr())
				break
			}

			// if peer closed the connection
			if err == io.EOF {
				log.Printf("conn: peer disconnected %v: %v", c.ID, c.RemoteAddr())
				break
			}

			log.Printf("conn: error while receiving message: %v", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			c.wg.Add(1)
			go func() {
				defer recoverAndLog(c, &c.wg)
				if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
					log.Printf("conn: error while handling request: %v", err)
					c.Close()
				}
			}()

			continue
		}

		// if the message is not a JSON-RPC message
		if m.ID == "" || (m.Result == nil && m.Error == nil) {
			log.Printf("conn: received an unknown message %v: %v, %v", c.ID, c.RemoteAddr(), m)
			break
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			c.wg.Add(1)
			go func() {
				defer recoverAndLog(c, &c.wg)
				err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
				c.resRoutes.Delete(m.ID)
				if err != nil {
					log.Printf("conn: error while handling response: %v", err)
					c.Close()
				}
			}()
		} else {
			log.Printf("conn: error while handling response: got response to a request with unknown ID: %v", m.ID)
			break
		}
	}
}

func recoverAndLog(c *Conn, wg *sync.WaitGroup) {
	defer wg.Done()
	if err := recover(); err != nil {
		const size = 64 << 10
		buf := make([]byte, size)
		buf = buf[:runtime.Stack(buf, false)]
		log.Printf("conn: panic handling response %v: %v\n%s", c.RemoteAddr(), err, buf)
	}
}
<|endoftext|>


<commit_before>package gosocks5

import (
	"io"
	"log"
	"net"
	"sync"
	"time"
)

type Config struct {
	Methods        []uint8
	SelectMethod   func(methods ...uint8) uint8
	MethodSelected func(method uint8, conn net.Conn) (net.Conn, error)
}

func defaultConfig() *Config {
	return &Config{}
}

type Conn struct {
	c              net.Conn
	config         *Config
	method         uint8
	isClient       bool
	handshaked     bool
	handshakeMutex sync.Mutex
	handshakeErr   error
}

func ClientConn(conn net.Conn, config *Config) *Conn {
	return &Conn{
		c:        conn,
		config:   config,
		isClient: true,
	}
}

func ServerConn(conn net.Conn, config *Config) *Conn {
	return &Conn{
		c:      conn,
		config: config,
	}
}

func (conn *Conn) Handleshake() error {
	conn.handshakeMutex.Lock()
	defer conn.handshakeMutex.Unlock()

	if err := conn.handshakeErr; err != nil {
		return err
	}
	if conn.handshaked {
		return nil
	}

	if conn.isClient {
		conn.handshakeErr = conn.clientHandshake()
	} else {
		conn.handshakeErr = conn.serverHandshake()
	}

	return conn.handshakeErr
}

func (conn *Conn) clientHandshake() error {
	if conn.config == nil {
		conn.config = defaultConfig()
	}

	nm := len(conn.config.Methods)
	if nm == 0 {
		nm = 1
	}

	b := make([]byte, 2+nm)
	b[0] = Ver5
	b[1] = uint8(nm)
	copy(b[2:], conn.config.Methods)

	if _, err := conn.c.Write(b); err != nil {
		return err
	}

	if _, err := io.ReadFull(conn.c, b[:2]); err != nil {
		return err
	}

	if b[0] != Ver5 {
		return ErrBadVersion
	}

	if conn.config.MethodSelected != nil {
		c, err := conn.config.MethodSelected(b[1], conn.c)
		if err != nil {
			return err
		}
		conn.c = c
	}
	conn.method = b[1]
	log.Println("method:", conn.method)
	conn.handshaked = true
	return nil
}

func (conn *Conn) serverHandshake() error {
	if conn.config == nil {
		conn.config = defaultConfig()
	}

	methods, err := ReadMethods(conn.c)
	if err != nil {
		return err
	}

	method := MethodNoAuth
	if conn.config.SelectMethod != nil {
		method = conn.config.SelectMethod(methods...)
	}

	if _, err := conn.c.Write([]byte{Ver5, method}); err != nil {
		return err
	}

	if conn.config.MethodSelected != nil {
		c, err := conn.config.MethodSelected(method, conn.c)
		if err != nil {
			return err
		}
		conn.c = c
	}
	conn.method = method
	log.Println("method:", method)
	conn.handshaked = true
	return nil
}

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

func (conn *Conn) Write(b []byte) (n int, err error) {
	if err = conn.Handleshake(); err != nil {
		return
	}
	return conn.c.Write(b)
}

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

func (conn *Conn) LocalAddr() net.Addr {
	return conn.c.LocalAddr()
}

func (conn *Conn) RemoteAddr() net.Addr {
	return conn.c.RemoteAddr()
}

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

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

func (conn *Conn) SetWriteDeadline(t time.Time) error {
	return conn.c.SetWriteDeadline(t)
}
<commit_msg>rm print<commit_after>package gosocks5

import (
	"io"
	//"log"
	"net"
	"sync"
	"time"
)

type Config struct {
	Methods        []uint8
	SelectMethod   func(methods ...uint8) uint8
	MethodSelected func(method uint8, conn net.Conn) (net.Conn, error)
}

func defaultConfig() *Config {
	return &Config{}
}

type Conn struct {
	c              net.Conn
	config         *Config
	method         uint8
	isClient       bool
	handshaked     bool
	handshakeMutex sync.Mutex
	handshakeErr   error
}

func ClientConn(conn net.Conn, config *Config) *Conn {
	return &Conn{
		c:        conn,
		config:   config,
		isClient: true,
	}
}

func ServerConn(conn net.Conn, config *Config) *Conn {
	return &Conn{
		c:      conn,
		config: config,
	}
}

func (conn *Conn) Handleshake() error {
	conn.handshakeMutex.Lock()
	defer conn.handshakeMutex.Unlock()

	if err := conn.handshakeErr; err != nil {
		return err
	}
	if conn.handshaked {
		return nil
	}

	if conn.isClient {
		conn.handshakeErr = conn.clientHandshake()
	} else {
		conn.handshakeErr = conn.serverHandshake()
	}

	return conn.handshakeErr
}

func (conn *Conn) clientHandshake() error {
	if conn.config == nil {
		conn.config = defaultConfig()
	}

	nm := len(conn.config.Methods)
	if nm == 0 {
		nm = 1
	}

	b := make([]byte, 2+nm)
	b[0] = Ver5
	b[1] = uint8(nm)
	copy(b[2:], conn.config.Methods)

	if _, err := conn.c.Write(b); err != nil {
		return err
	}

	if _, err := io.ReadFull(conn.c, b[:2]); err != nil {
		return err
	}

	if b[0] != Ver5 {
		return ErrBadVersion
	}

	if conn.config.MethodSelected != nil {
		c, err := conn.config.MethodSelected(b[1], conn.c)
		if err != nil {
			return err
		}
		conn.c = c
	}
	conn.method = b[1]
	//log.Println("method:", conn.method)
	conn.handshaked = true
	return nil
}

func (conn *Conn) serverHandshake() error {
	if conn.config == nil {
		conn.config = defaultConfig()
	}

	methods, err := ReadMethods(conn.c)
	if err != nil {
		return err
	}

	method := MethodNoAuth
	if conn.config.SelectMethod != nil {
		method = conn.config.SelectMethod(methods...)
	}

	if _, err := conn.c.Write([]byte{Ver5, method}); err != nil {
		return err
	}

	if conn.config.MethodSelected != nil {
		c, err := conn.config.MethodSelected(method, conn.c)
		if err != nil {
			return err
		}
		conn.c = c
	}
	conn.method = method
	//log.Println("method:", method)
	conn.handshaked = true
	return nil
}

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

func (conn *Conn) Write(b []byte) (n int, err error) {
	if err = conn.Handleshake(); err != nil {
		return
	}
	return conn.c.Write(b)
}

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

func (conn *Conn) LocalAddr() net.Addr {
	return conn.c.LocalAddr()
}

func (conn *Conn) RemoteAddr() net.Addr {
	return conn.c.RemoteAddr()
}

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

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

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


<commit_before>package main

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"net/textproto"
	"strings"
	"sync"
	"time"
)

type connType uint32

const (
	connWhisperConn = iota
	connReadConn
	connSendConn
	connDelete
)

type connection struct {
	sync.Mutex
	conn     net.Conn
	active   bool
	anon     bool
	joins    []string
	msgCount int
	lastUse  time.Time
	alive    bool
	conntype connType
	bot      *bot
	name     string
}

func newConnection(t connType) *connection {
	return &connection{
		joins:    make([]string, 0),
		conntype: t,
		lastUse:  time.Now(),
		name:     randomHash(),
	}
}

func (conn *connection) login(pass string, nick string) {
	conn.anon = pass == ""
	if !conn.anon {
		conn.send("PASS " + pass)
		conn.send("NICK " + nick)
		return
	}
	conn.send("NICK justinfan123")
}

func (conn *connection) close() {
	if conn.conn != nil {
		conn.conn.Close()
	}
	for _, channel := range conn.joins {
		conn.part(channel)
	}
	conn.alive = false
}

func (conn *connection) part(channel string) {
	channel = strings.ToLower(channel)
	for i, ch := range conn.joins {
		if ch == channel {
			conn.joins = append(conn.joins[:i], conn.joins[i+1:]...)
		}
	}
}

func (conn *connection) restore() {
	defer func() {
		if r := recover(); r != nil {
			Log.Error("cannot restore connection")
		}
	}()
	if conn.conntype == connReadConn {
		var i int
		var channels []string
		for index, co := range conn.bot.readconns {
			if conn == co {
				i = index
				channels = co.joins
				break
			}
		}
		Log.Error("readconn died, lost joins:", channels)
		conn.bot.Lock()
		conn.bot.readconns = append(conn.bot.readconns[:i], conn.bot.readconns[i+1:]...)
		conn.bot.Unlock()
		for _, channel := range channels {
			conns := conn.bot.channels[channel]
			for i, co := range conns {
				if conn == co {
					conn.bot.Lock()
					conn.bot.channels[channel] = append(conns[:i], conns[i+1:]...)
					conn.bot.Unlock()
					conn.part(channel)
				}
			}
			conn.bot.join <- channel

		}

	} else if conn.conntype == connSendConn {
		Log.Error("sendconn died")
		var i int
		for index, co := range conn.bot.sendconns {
			if conn == co {
				i = index
				break
			}
		}
		conn.bot.Lock()
		conn.bot.sendconns = append(conn.bot.sendconns[:i], conn.bot.sendconns[i+1:]...)
		conn.bot.Unlock()
	} else if conn.conntype == connWhisperConn {
		Log.Error("whisperconn died, reconnecting")
		conn.close()
		conn.bot.newConn(connWhisperConn)
	}
	conn.conntype = connDelete
}

func (conn *connection) connect(client *Client, pass string, nick string) {
	dialer := &net.Dialer{
		KeepAlive: time.Second * 10,
	}

	conn.bot = client.bot
	c, err := tls.DialWithDialer(dialer, "tcp", *addr, &tls.Config{})
	if err != nil {
		Log.Error("unable to connect to irc server", err)
		time.Sleep(2 * time.Second)
		conn.restore()
		return
	}
	conn.conn = c

	conn.send("CAP REQ :twitch.tv/tags twitch.tv/commands")
	conn.login(pass, nick)

	defer func() {
		if r := recover(); r != nil {
			Log.Error("error connecting")
		}
		conn.restore()
	}()
	tp := textproto.NewReader(bufio.NewReader(conn.conn))

	for {
		line, err := tp.ReadLine()
		if err != nil {
			Log.Fatalf("[READERROR:%s] %s", conn.name, err.Error())
			// conn.restore()
			return
		}
		Log.Debugf("[TWITCH:%s] %s", conn.name, line)
		if conn.conntype == connDelete {
			conn.restore()
		}
		if strings.HasPrefix(line, "PING") {
			conn.send(strings.Replace(line, "PING", "PONG", 1))
		} else if strings.HasPrefix(line, "PONG") {
			Log.Debug("PONG")
		} else {
			if isWhisper(line) && conn.conntype != connWhisperConn {
				// throw away message
			} else {
				client.toClient <- line
			}
		}
		conn.active = true
	}
}

func isWhisper(line string) bool {
	if !strings.Contains(line, ".tmi.twitch.tv WHISPER ") {
		return false
	}
	spl := strings.SplitN(line, " :", 3)
	if strings.Contains(spl[1], ".tmi.twitch.tv WHISPER ") {
		return true
	}
	return false
}

func (conn *connection) send(msg string) error {
	if conn.conn == nil {
		Log.Error("conn is nil", conn, conn.conn)
		return errors.New("connection is nil")
	}
	_, err := fmt.Fprint(conn.conn, msg+"\r\n")
	if err != nil {
		Log.Error("error sending message")
		return err
	}
	Log.Debugf("[OUTGOING:%s] %s", conn.name, msg)
	return nil
}

func (conn *connection) reduceMsgCount() {
	conn.msgCount--
}

func (conn *connection) countMsg() {
	conn.msgCount++
	time.AfterFunc(30*time.Second, conn.reduceMsgCount)
}

func randomHash() string {
	n := 5
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return fmt.Sprintf("%X", b)
}
<commit_msg>readd restore<commit_after>package main

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"net/textproto"
	"strings"
	"sync"
	"time"
)

type connType uint32

const (
	connWhisperConn = iota
	connReadConn
	connSendConn
	connDelete
)

type connection struct {
	sync.Mutex
	conn     net.Conn
	active   bool
	anon     bool
	joins    []string
	msgCount int
	lastUse  time.Time
	alive    bool
	conntype connType
	bot      *bot
	name     string
}

func newConnection(t connType) *connection {
	return &connection{
		joins:    make([]string, 0),
		conntype: t,
		lastUse:  time.Now(),
		name:     randomHash(),
	}
}

func (conn *connection) login(pass string, nick string) {
	conn.anon = pass == ""
	if !conn.anon {
		conn.send("PASS " + pass)
		conn.send("NICK " + nick)
		return
	}
	conn.send("NICK justinfan123")
}

func (conn *connection) close() {
	if conn.conn != nil {
		conn.conn.Close()
	}
	for _, channel := range conn.joins {
		conn.part(channel)
	}
	conn.alive = false
}

func (conn *connection) part(channel string) {
	channel = strings.ToLower(channel)
	for i, ch := range conn.joins {
		if ch == channel {
			conn.joins = append(conn.joins[:i], conn.joins[i+1:]...)
		}
	}
}

func (conn *connection) restore() {
	defer func() {
		if r := recover(); r != nil {
			Log.Error("cannot restore connection")
		}
	}()
	if conn.conntype == connReadConn {
		var i int
		var channels []string
		for index, co := range conn.bot.readconns {
			if conn == co {
				i = index
				channels = co.joins
				break
			}
		}
		Log.Error("readconn died, lost joins:", channels)
		conn.bot.Lock()
		conn.bot.readconns = append(conn.bot.readconns[:i], conn.bot.readconns[i+1:]...)
		conn.bot.Unlock()
		for _, channel := range channels {
			conns := conn.bot.channels[channel]
			for i, co := range conns {
				if conn == co {
					conn.bot.Lock()
					conn.bot.channels[channel] = append(conns[:i], conns[i+1:]...)
					conn.bot.Unlock()
					conn.part(channel)
				}
			}
			conn.bot.join <- channel

		}

	} else if conn.conntype == connSendConn {
		Log.Error("sendconn died")
		var i int
		for index, co := range conn.bot.sendconns {
			if conn == co {
				i = index
				break
			}
		}
		conn.bot.Lock()
		conn.bot.sendconns = append(conn.bot.sendconns[:i], conn.bot.sendconns[i+1:]...)
		conn.bot.Unlock()
	} else if conn.conntype == connWhisperConn {
		Log.Error("whisperconn died, reconnecting")
		conn.close()
		conn.bot.newConn(connWhisperConn)
	}
	conn.conntype = connDelete
}

func (conn *connection) connect(client *Client, pass string, nick string) {
	dialer := &net.Dialer{
		KeepAlive: time.Second * 10,
	}

	conn.bot = client.bot
	c, err := tls.DialWithDialer(dialer, "tcp", *addr, &tls.Config{})
	if err != nil {
		Log.Error("unable to connect to irc server", err)
		time.Sleep(2 * time.Second)
		conn.restore()
		return
	}
	conn.conn = c

	conn.send("CAP REQ :twitch.tv/tags twitch.tv/commands")
	conn.login(pass, nick)

	defer func() {
		if r := recover(); r != nil {
			Log.Error("error connecting")
		}
		conn.restore()
	}()
	tp := textproto.NewReader(bufio.NewReader(conn.conn))

	for {
		line, err := tp.ReadLine()
		if err != nil {
			Log.Errorf("[READERROR:%s] %s", conn.name, err.Error())
			conn.restore()
			return
		}
		Log.Debugf("[TWITCH:%s] %s", conn.name, line)
		if conn.conntype == connDelete {
			conn.restore()
		}
		if strings.HasPrefix(line, "PING") {
			conn.send(strings.Replace(line, "PING", "PONG", 1))
		} else if strings.HasPrefix(line, "PONG") {
			Log.Debug("PONG")
		} else {
			if isWhisper(line) && conn.conntype != connWhisperConn {
				// throw away message
			} else {
				client.toClient <- line
			}
		}
		conn.active = true
	}
}

func isWhisper(line string) bool {
	if !strings.Contains(line, ".tmi.twitch.tv WHISPER ") {
		return false
	}
	spl := strings.SplitN(line, " :", 3)
	if strings.Contains(spl[1], ".tmi.twitch.tv WHISPER ") {
		return true
	}
	return false
}

func (conn *connection) send(msg string) error {
	if conn.conn == nil {
		Log.Error("conn is nil", conn, conn.conn)
		return errors.New("connection is nil")
	}
	_, err := fmt.Fprint(conn.conn, msg+"\r\n")
	if err != nil {
		Log.Error("error sending message")
		return err
	}
	Log.Debugf("[OUTGOING:%s] %s", conn.name, msg)
	return nil
}

func (conn *connection) reduceMsgCount() {
	conn.msgCount--
}

func (conn *connection) countMsg() {
	conn.msgCount++
	time.AfterFunc(30*time.Second, conn.reduceMsgCount)
}

func randomHash() string {
	n := 5
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return fmt.Sprintf("%X", b)
}
<|endoftext|>


<commit_before>package gotcp

import (
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Error type
var (
	ErrConnClosing   = errors.New("use of closed network connection")
	ErrWriteBlocking = errors.New("write packet was blocking")
	ErrReadBlocking  = errors.New("read packet was blocking")
)

// Conn exposes a set of callbacks for the various events that occur on a connection
type Conn struct {
	srv               *Server
	conn              *net.TCPConn  // the raw connection
	extraData         interface{}   // to save extra data
	closeOnce         sync.Once     // close the conn, once, per instance
	closeFlag         int32         // close flag
	closeChan         chan struct{} // close chanel
	packetSendChan    chan Packet   // packet send chanel
	packetReceiveChan chan Packet   // packeet receive chanel
}

// ConnCallback is an interface of methods that are used as callbacks on a connection
type ConnCallback interface {
	// OnConnect is called when the connection was accepted,
	// If the return value of false is closed
	OnConnect(*Conn) bool

	// OnMessage is called when the connection receives a packet,
	// If the return value of false is closed
	OnMessage(*Conn, Packet) bool

	// OnClose is called when the connection closed
	OnClose(*Conn)
}

// newConn returns a wrapper of raw conn
func newConn(conn *net.TCPConn, srv *Server) *Conn {
	return &Conn{
		srv:               srv,
		conn:              conn,
		closeChan:         make(chan struct{}),
		packetSendChan:    make(chan Packet, srv.config.PacketSendChanLimit),
		packetReceiveChan: make(chan Packet, srv.config.PacketReceiveChanLimit),
	}
}

// GetExtraData gets the extra data from the Conn
func (c *Conn) GetExtraData() interface{} {
	return c.extraData
}

// PutExtraData puts the extra data with the Conn
func (c *Conn) PutExtraData(data interface{}) {
	c.extraData = data
}

// GetRawConn returns the raw net.TCPConn from the Conn
func (c *Conn) GetRawConn() *net.TCPConn {
	return c.conn
}

// Close closes the connection
func (c *Conn) Close() {
	c.closeOnce.Do(func() {
		atomic.StoreInt32(&c.closeFlag, 1)
		close(c.closeChan)
		c.conn.Close()
		c.srv.callback.OnClose(c)
	})
}

// IsClosed indicates whether or not the connection is closed
func (c *Conn) IsClosed() bool {
	return atomic.LoadInt32(&c.closeFlag) == 1
}

// AsyncReadPacket async reads a packet, this method will never block
func (c *Conn) AsyncReadPacket(timeout time.Duration) (Packet, error) {
	if c.IsClosed() {
		return nil, ErrConnClosing
	}

	if timeout == 0 {
		select {
		case p := <-c.packetReceiveChan:
			return p, nil

		default:
			return nil, ErrReadBlocking
		}

	} else {
		select {
		case p := <-c.packetReceiveChan:
			return p, nil

		case <-c.closeChan:
			return nil, ErrConnClosing

		case <-time.After(timeout):
			return nil, ErrReadBlocking
		}
	}
}

// AsyncWritePacket async writes a packet, this method will never block
func (c *Conn) AsyncWritePacket(p Packet, timeout time.Duration) error {
	if c.IsClosed() {
		return ErrConnClosing
	}

	if timeout == 0 {
		select {
		case c.packetSendChan <- p:
			return nil

		default:
			return ErrWriteBlocking
		}

	} else {
		select {
		case c.packetSendChan <- p:
			return nil

		case <-c.closeChan:
			return ErrConnClosing

		case <-time.After(timeout):
			return ErrWriteBlocking
		}
	}
}

// Do it
func (c *Conn) Do() {
	if !c.srv.callback.OnConnect(c) {
		return
	}

	go c.handleLoop()
	go c.readLoop()
	go c.writeLoop()
}

func (c *Conn) readLoop() {
	c.srv.waitGroup.Add(1)
	defer func() {
		recover()
		c.Close()
		c.srv.waitGroup.Done()
	}()

	for {
		select {
		case <-c.srv.exitChan:
			return

		case <-c.closeChan:
			return

		default:
		}

		p, err := c.srv.protocol.ReadPacket(c.conn)
		if err != nil {
			return
		}

		c.packetReceiveChan <- p
	}
}

func (c *Conn) writeLoop() {
	c.srv.waitGroup.Add(1)
	defer func() {
		recover()
		c.Close()
		c.srv.waitGroup.Done()
	}()

	for {
		select {
		case <-c.srv.exitChan:
			return

		case <-c.closeChan:
			return

		case p := <-c.packetSendChan:
			if _, err := c.conn.Write(p.Serialize()); err != nil {
				return
			}
		}
	}
}

func (c *Conn) handleLoop() {
	c.srv.waitGroup.Add(1)
	defer func() {
		recover()
		c.Close()
		c.srv.waitGroup.Done()
	}()

	for {
		select {
		case <-c.srv.exitChan:
			return

		case <-c.closeChan:
			return

		case p := <-c.packetReceiveChan:
			if !c.srv.callback.OnMessage(c, p) {
				return
			}
		}
	}
}
<commit_msg>fix possible panic and blocking<commit_after>package gotcp

import (
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Error type
var (
	ErrConnClosing   = errors.New("use of closed network connection")
	ErrWriteBlocking = errors.New("write packet was blocking")
	ErrReadBlocking  = errors.New("read packet was blocking")
)

// Conn exposes a set of callbacks for the various events that occur on a connection
type Conn struct {
	srv               *Server
	conn              *net.TCPConn  // the raw connection
	extraData         interface{}   // to save extra data
	closeOnce         sync.Once     // close the conn, once, per instance
	closeFlag         int32         // close flag
	closeChan         chan struct{} // close chanel
	packetSendChan    chan Packet   // packet send chanel
	packetReceiveChan chan Packet   // packeet receive chanel
}

// ConnCallback is an interface of methods that are used as callbacks on a connection
type ConnCallback interface {
	// OnConnect is called when the connection was accepted,
	// If the return value of false is closed
	OnConnect(*Conn) bool

	// OnMessage is called when the connection receives a packet,
	// If the return value of false is closed
	OnMessage(*Conn, Packet) bool

	// OnClose is called when the connection closed
	OnClose(*Conn)
}

// newConn returns a wrapper of raw conn
func newConn(conn *net.TCPConn, srv *Server) *Conn {
	return &Conn{
		srv:               srv,
		conn:              conn,
		closeChan:         make(chan struct{}),
		packetSendChan:    make(chan Packet, srv.config.PacketSendChanLimit),
		packetReceiveChan: make(chan Packet, srv.config.PacketReceiveChanLimit),
	}
}

// GetExtraData gets the extra data from the Conn
func (c *Conn) GetExtraData() interface{} {
	return c.extraData
}

// PutExtraData puts the extra data with the Conn
func (c *Conn) PutExtraData(data interface{}) {
	c.extraData = data
}

// GetRawConn returns the raw net.TCPConn from the Conn
func (c *Conn) GetRawConn() *net.TCPConn {
	return c.conn
}

// Close closes the connection
func (c *Conn) Close() {
	c.closeOnce.Do(func() {
		atomic.StoreInt32(&c.closeFlag, 1)
		close(c.closeChan)
		close(c.packetSendChan)
		close(c.packetReceiveChan)
		c.conn.Close()
		c.srv.callback.OnClose(c)
	})
}

// IsClosed indicates whether or not the connection is closed
func (c *Conn) IsClosed() bool {
	return atomic.LoadInt32(&c.closeFlag) == 1
}

// AsyncWritePacket async writes a packet, this method will never block
func (c *Conn) AsyncWritePacket(p Packet, timeout time.Duration) (err error) {
	if c.IsClosed() {
		return ErrConnClosing
	}

	defer func() {
		if e := recover(); e != nil {
			err = ErrConnClosing
		}
	}()

	if timeout == 0 {
		select {
		case c.packetSendChan <- p:
			return nil

		default:
			return ErrWriteBlocking
		}

	} else {
		select {
		case c.packetSendChan <- p:
			return nil

		case <-c.closeChan:
			return ErrConnClosing

		case <-time.After(timeout):
			return ErrWriteBlocking
		}
	}
}

// Do it
func (c *Conn) Do() {
	if !c.srv.callback.OnConnect(c) {
		return
	}

	go c.handleLoop()
	go c.readLoop()
	go c.writeLoop()
}

func (c *Conn) readLoop() {
	c.srv.waitGroup.Add(1)
	defer func() {
		recover()
		c.Close()
		c.srv.waitGroup.Done()
	}()

	for {
		select {
		case <-c.srv.exitChan:
			return

		case <-c.closeChan:
			return

		default:
		}

		p, err := c.srv.protocol.ReadPacket(c.conn)
		if err != nil {
			return
		}

		c.packetReceiveChan <- p
	}
}

func (c *Conn) writeLoop() {
	c.srv.waitGroup.Add(1)
	defer func() {
		recover()
		c.Close()
		c.srv.waitGroup.Done()
	}()

	for {
		select {
		case <-c.srv.exitChan:
			return

		case <-c.closeChan:
			return

		case p := <-c.packetSendChan:
			if c.IsClosed() {
				return
			}
			if _, err := c.conn.Write(p.Serialize()); err != nil {
				return
			}
		}
	}
}

func (c *Conn) handleLoop() {
	c.srv.waitGroup.Add(1)
	defer func() {
		recover()
		c.Close()
		c.srv.waitGroup.Done()
	}()

	for {
		select {
		case <-c.srv.exitChan:
			return

		case <-c.closeChan:
			return

		case p := <-c.packetReceiveChan:
			if c.IsClosed() {
				return
			}
			if !c.srv.callback.OnMessage(c, p) {
				return
			}
		}
	}
}
<|endoftext|>


<commit_before>package testblas

import (
	"math"
	"testing"
)

type Dgerer interface {
	Dger(m, n int, alpha float64, x []float64, incX int, y []float64, incY int, a []float64, lda int)
}

func DgerTest(t *testing.T, blasser Dgerer) {
	for _, test := range []struct {
		name string
		a    [][]float64
		m    int
		n    int
		x    []float64
		y    []float64
		incX int
		incY int

		trueAns [][]float64
	}{
		{
			name: "M gt N inc 1",
			m:    5,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
				{8, 9, -10},
				{-12, -14, -6},
			},
			x:       []float64{-2, -3, 0, 1, 2},
			y:       []float64{-1.1, 5, 0},
			incX:    1,
			incY:    1,
			trueAns: [][]float64{{3.5, -7.6, 3.5}, {5.9, -12.2, 3.3}, {-1.3, -4.3, -9.7}, {6.9, 14, -10}, {-14.2, -4, -6}},
		},
		{
			name: "M eq N inc 1",
			m:    3,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
			},
			x:       []float64{-2, -3, 0},
			y:       []float64{-1.1, 5, 0},
			incX:    1,
			incY:    1,
			trueAns: [][]float64{{3.5, -7.6, 3.5}, {5.9, -12.2, 3.3}, {-1.3, -4.3, -9.7}},
		},

		{
			name: "M lt N inc 1",
			m:    3,
			n:    6,
			a: [][]float64{
				{1.3, 2.4, 3.5, 4.8, 1.11, -9},
				{2.6, 2.8, 3.3, -3.4, 6.2, -8.7},
				{-1.3, -4.3, -9.7, -3.1, 8.9, 8.9},
			},
			x:       []float64{-2, -3, 0},
			y:       []float64{-1.1, 5, 0, 9, 19, 22},
			incX:    1,
			incY:    1,
			trueAns: [][]float64{{3.5, -7.6, 3.5, -13.2, -36.89, -53}, {5.9, -12.2, 3.3, -30.4, -50.8, -74.7}, {-1.3, -4.3, -9.7, -3.1, 8.9, 8.9}},
		},
		{
			name: "M gt N inc not 1",
			m:    5,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
				{8, 9, -10},
				{-12, -14, -6},
			},
			x:       []float64{-2, -3, 0, 1, 2, 6, 0, 9, 7},
			y:       []float64{-1.1, 5, 0, 8, 7, -5, 7},
			incX:    2,
			incY:    3,
			trueAns: [][]float64{{3.5, -13.6, -10.5}, {2.6, 2.8, 3.3}, {-3.5, 11.7, 4.3}, {8, 9, -10}, {-19.700000000000003, 42, 43}},
		},
		{
			name: "M eq N inc not 1",
			m:    3,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
			},
			x:       []float64{-2, -3, 0, 8, 7, -9, 7, -6, 12, 6, 6, 6, -11},
			y:       []float64{-1.1, 5, 0, 0, 9, 8, 6},
			incX:    4,
			incY:    3,
			trueAns: [][]float64{{3.5, 2.4, -8.5}, {-5.1, 2.8, 45.3}, {-14.5, -4.3, 62.3}},
		},
		{
			name: "M lt N inc not 1",
			m:    3,
			n:    6,
			a: [][]float64{
				{1.3, 2.4, 3.5, 4.8, 1.11, -9},
				{2.6, 2.8, 3.3, -3.4, 6.2, -8.7},
				{-1.3, -4.3, -9.7, -3.1, 8.9, 8.9},
			},
			x:       []float64{-2, -3, 0, 0, 8, 0, 9, -3},
			y:       []float64{-1.1, 5, 0, 9, 19, 22, 11, -8.11, -9.22, 9.87, 7},
			incX:    3,
			incY:    2,
			trueAns: [][]float64{{3.5, 2.4, -34.5, -17.2, 19.55, -23}, {2.6, 2.8, 3.3, -3.4, 6.2, -8.7}, {-11.2, -4.3, 161.3, 95.9, -74.08, 71.9}},
		},
		{
			name:    "Y NaN element",
			m:       1,
			n:       1,
			a:       [][]float64{{1.3}},
			x:       []float64{1.3},
			y:       []float64{math.NaN()},
			incX:    1,
			incY:    1,
			trueAns: [][]float64{{math.NaN()}},
		},
	} {
		// TODO: Add tests where a is longer
		// TODO: Add panic tests
		// TODO: Add negative increment tests

		x := sliceCopy(test.x)
		y := sliceCopy(test.y)

		a := sliceOfSliceCopy(test.a)

		// Test with row major
		alpha := 1.0
		aFlat := flatten(a)
		blasser.Dger(test.m, test.n, alpha, x, test.incX, y, test.incY, aFlat, test.n)
		ans := unflatten(aFlat, test.m, test.n)
		dgercomp(t, x, test.x, y, test.y, ans, test.trueAns, test.name+" row maj")

		// Test with different alpha
		alpha = 4.0
		aFlat = flatten(a)
		blasser.Dger(test.m, test.n, alpha, x, test.incX, y, test.incY, aFlat, test.n)
		ans = unflatten(aFlat, test.m, test.n)
		trueCopy := sliceOfSliceCopy(test.trueAns)
		for i := range trueCopy {
			for j := range trueCopy[i] {
				trueCopy[i][j] = alpha*(trueCopy[i][j]-a[i][j]) + a[i][j]
			}
		}
		dgercomp(t, x, test.x, y, test.y, ans, trueCopy, test.name+" row maj alpha")
	}
}

func dgercomp(t *testing.T, x, xCopy, y, yCopy []float64, ans [][]float64, trueAns [][]float64, name string) {
	if !dSliceEqual(x, xCopy) {
		t.Errorf("case %v: x modified during call to dger\n%v\n%v", name, x, xCopy)
	}
	if !dSliceEqual(y, yCopy) {
		t.Errorf("case %v: y modified during call to dger\n%v\n%v", name, y, yCopy)
	}

	for i := range ans {
		if !dSliceTolEqual(ans[i], trueAns[i]) {
			t.Errorf("case %v: answer mismatch at %v. Expected %v, Found %v", name, i, trueAns, ans)
			break
		}
	}
}
<commit_msg>blas/testblas: Added 7x7 tests to dger.<commit_after>package testblas

import (
	"math"
	"testing"
)

type Dgerer interface {
	Dger(m, n int, alpha float64, x []float64, incX int, y []float64, incY int, a []float64, lda int)
}

func DgerTest(t *testing.T, blasser Dgerer) {
	for _, test := range []struct {
		name string
		a    [][]float64
		m    int
		n    int
		x    []float64
		y    []float64
		incX int
		incY int

		want [][]float64
	}{
		{
			name: "M gt N inc 1",
			m:    5,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
				{8, 9, -10},
				{-12, -14, -6},
			},
			x:    []float64{-2, -3, 0, 1, 2},
			y:    []float64{-1.1, 5, 0},
			incX: 1,
			incY: 1,
			want: [][]float64{{3.5, -7.6, 3.5}, {5.9, -12.2, 3.3}, {-1.3, -4.3, -9.7}, {6.9, 14, -10}, {-14.2, -4, -6}},
		},
		{
			name: "M eq N inc 1",
			m:    3,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
			},
			x:    []float64{-2, -3, 0},
			y:    []float64{-1.1, 5, 0},
			incX: 1,
			incY: 1,
			want: [][]float64{{3.5, -7.6, 3.5}, {5.9, -12.2, 3.3}, {-1.3, -4.3, -9.7}},
		},

		{
			name: "M lt N inc 1",
			m:    3,
			n:    6,
			a: [][]float64{
				{1.3, 2.4, 3.5, 4.8, 1.11, -9},
				{2.6, 2.8, 3.3, -3.4, 6.2, -8.7},
				{-1.3, -4.3, -9.7, -3.1, 8.9, 8.9},
			},
			x:    []float64{-2, -3, 0},
			y:    []float64{-1.1, 5, 0, 9, 19, 22},
			incX: 1,
			incY: 1,
			want: [][]float64{{3.5, -7.6, 3.5, -13.2, -36.89, -53}, {5.9, -12.2, 3.3, -30.4, -50.8, -74.7}, {-1.3, -4.3, -9.7, -3.1, 8.9, 8.9}},
		},
		{
			name: "M gt N inc not 1",
			m:    5,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
				{8, 9, -10},
				{-12, -14, -6},
			},
			x:    []float64{-2, -3, 0, 1, 2, 6, 0, 9, 7},
			y:    []float64{-1.1, 5, 0, 8, 7, -5, 7},
			incX: 2,
			incY: 3,
			want: [][]float64{{3.5, -13.6, -10.5}, {2.6, 2.8, 3.3}, {-3.5, 11.7, 4.3}, {8, 9, -10}, {-19.700000000000003, 42, 43}},
		},
		{
			name: "M eq N inc not 1",
			m:    3,
			n:    3,
			a: [][]float64{
				{1.3, 2.4, 3.5},
				{2.6, 2.8, 3.3},
				{-1.3, -4.3, -9.7},
			},
			x:    []float64{-2, -3, 0, 8, 7, -9, 7, -6, 12, 6, 6, 6, -11},
			y:    []float64{-1.1, 5, 0, 0, 9, 8, 6},
			incX: 4,
			incY: 3,
			want: [][]float64{{3.5, 2.4, -8.5}, {-5.1, 2.8, 45.3}, {-14.5, -4.3, 62.3}},
		},
		{
			name: "M lt N inc not 1",
			m:    3,
			n:    6,
			a: [][]float64{
				{1.3, 2.4, 3.5, 4.8, 1.11, -9},
				{2.6, 2.8, 3.3, -3.4, 6.2, -8.7},
				{-1.3, -4.3, -9.7, -3.1, 8.9, 8.9},
			},
			x:    []float64{-2, -3, 0, 0, 8, 0, 9, -3},
			y:    []float64{-1.1, 5, 0, 9, 19, 22, 11, -8.11, -9.22, 9.87, 7},
			incX: 3,
			incY: 2,
			want: [][]float64{{3.5, 2.4, -34.5, -17.2, 19.55, -23}, {2.6, 2.8, 3.3, -3.4, 6.2, -8.7}, {-11.2, -4.3, 161.3, 95.9, -74.08, 71.9}},
		},
		{
			name: "Y NaN element",
			m:    1,
			n:    1,
			a:    [][]float64{{1.3}},
			x:    []float64{1.3},
			y:    []float64{math.NaN()},
			incX: 1,
			incY: 1,
			want: [][]float64{{math.NaN()}},
		},
		{
			name: "M eq N large inc 1",
			m:    7,
			n:    7,
			x:    []float64{6.2, -5, 88.68, 43.4, -30.5, -40.2, 19.9},
			y:    []float64{1.5, 21.7, -28.7, -11.9, 18.1, 3.1, 21},
			a: [][]float64{
				{-20.5, 17.1, -8.4, -23.8, 3.9, 7.7, 6.25},
				{2.9, -0.29, 25.6, -9.4, 36.5, 9.7, 2.3},
				{4.1, -34.1, 10.3, 4.5, -42.05, 9.4, 4},
				{19.2, 9.8, -32.7, 4.1, 4.4, -22.5, -7.8},
				{3.6, -24.5, 21.7, 8.6, -13.82, 38.05, -2.29},
				{39.4, -40.5, 7.9, -2.5, -7.7, 18.1, -25.5},
				{-18.5, 43.2, 2.1, 30.1, 3.02, -31.1, -7.6},
			},
			incX: 1,
			incY: 1,
			want: [][]float64{
				{-11.2, 151.64, -186.34, -97.58, 116.12, 26.92, 136.45},
				{-4.6, -108.79, 169.1, 50.1, -54, -5.8, -102.7},
				{137.12, 1890.256, -2534.816, -1050.792, 1563.058, 284.308, 1866.28},
				{84.3, 951.58, -1278.28, -512.36, 789.94, 112.04, 903.6},
				{-42.15, -686.35, 897.05, 371.55, -565.87, -56.5, -642.79},
				{-20.9, -912.84, 1161.64, 475.88, -735.32, -106.52, -869.7},
				{11.35, 475.03, -569.03, -206.71, 363.21, 30.59, 410.3},
			},
		},
		{
			name: "M eq N large inc not 1",
			m:    7,
			n:    7,
			x:    []float64{6.2, 100, 200, -5, 300, 400, 88.68, 100, 200, 43.4, 300, 400, -30.5, 100, 200, -40.2, 300, 400, 19.9},
			y:    []float64{1.5, 100, 200, 300, 21.7, 100, 200, 300, -28.7, 100, 200, 300, -11.9, 100, 200, 300, 18.1, 100, 200, 300, 3.1, 100, 200, 300, 21},
			a: [][]float64{
				{-20.5, 17.1, -8.4, -23.8, 3.9, 7.7, 6.25},
				{2.9, -0.29, 25.6, -9.4, 36.5, 9.7, 2.3},
				{4.1, -34.1, 10.3, 4.5, -42.05, 9.4, 4},
				{19.2, 9.8, -32.7, 4.1, 4.4, -22.5, -7.8},
				{3.6, -24.5, 21.7, 8.6, -13.82, 38.05, -2.29},
				{39.4, -40.5, 7.9, -2.5, -7.7, 18.1, -25.5},
				{-18.5, 43.2, 2.1, 30.1, 3.02, -31.1, -7.6},
			},
			incX: 3,
			incY: 4,
			want: [][]float64{
				{-11.2, 151.64, -186.34, -97.58, 116.12, 26.92, 136.45},
				{-4.6, -108.79, 169.1, 50.1, -54, -5.8, -102.7},
				{137.12, 1890.256, -2534.816, -1050.792, 1563.058, 284.308, 1866.28},
				{84.3, 951.58, -1278.28, -512.36, 789.94, 112.04, 903.6},
				{-42.15, -686.35, 897.05, 371.55, -565.87, -56.5, -642.79},
				{-20.9, -912.84, 1161.64, 475.88, -735.32, -106.52, -869.7},
				{11.35, 475.03, -569.03, -206.71, 363.21, 30.59, 410.3},
			},
		},
	} {
		// TODO: Add tests where a is longer
		// TODO: Add panic tests
		// TODO: Add negative increment tests

		x := sliceCopy(test.x)
		y := sliceCopy(test.y)

		a := sliceOfSliceCopy(test.a)

		// Test with row major
		alpha := 1.0
		aFlat := flatten(a)
		blasser.Dger(test.m, test.n, alpha, x, test.incX, y, test.incY, aFlat, test.n)
		ans := unflatten(aFlat, test.m, test.n)
		dgercomp(t, x, test.x, y, test.y, ans, test.want, test.name+" row maj")

		// Test with different alpha
		alpha = 4.0
		aFlat = flatten(a)
		blasser.Dger(test.m, test.n, alpha, x, test.incX, y, test.incY, aFlat, test.n)
		ans = unflatten(aFlat, test.m, test.n)
		trueCopy := sliceOfSliceCopy(test.want)
		for i := range trueCopy {
			for j := range trueCopy[i] {
				trueCopy[i][j] = alpha*(trueCopy[i][j]-a[i][j]) + a[i][j]
			}
		}
		dgercomp(t, x, test.x, y, test.y, ans, trueCopy, test.name+" row maj alpha")
	}
}

func dgercomp(t *testing.T, x, xCopy, y, yCopy []float64, ans [][]float64, trueAns [][]float64, name string) {
	if !dSliceEqual(x, xCopy) {
		t.Errorf("case %v: x modified during call to dger\n%v\n%v", name, x, xCopy)
	}
	if !dSliceEqual(y, yCopy) {
		t.Errorf("case %v: y modified during call to dger\n%v\n%v", name, y, yCopy)
	}

	for i := range ans {
		if !dSliceTolEqual(ans[i], trueAns[i]) {
			t.Errorf("case %v: answer mismatch at %v.\nExpected %v,\nFound %v", name, i, trueAns, ans)
			break
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Qubit Digital Ltd.
// 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 logspray is a collection of tools for streaming and indexing
// large volumes of dynamic logs.

package indexer

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/QubitProducts/logspray/proto/logspray"
	"github.com/QubitProducts/logspray/ql"
	"github.com/gogo/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"github.com/pkg/errors"
)

// Search searches the shard file for messages in the provided time range,
// matched by matcher, and passes them to msgFunc. If the reverse is true the
// file will be searched in reverse order
func (s *ShardFile) Search(ctx context.Context, msgFunc logspray.MessageFunc, matcher ql.MatchFunc, from, to time.Time, count, offset uint64, reverse bool) error {
	s.RLock()
	defer s.RUnlock()

	r, err := os.Open(s.fn)
	defer r.Close()
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", s.fn)
	}

	hdr, err := readMessageFromFile(r)
	if err != nil {
		return errors.Wrapf(err, "failed to read file header %s", s.fn)
	}
	err = msgFunc(hdr)
	if err != nil {
		return err
	}
	if reverse {
		r.Seek(0, io.SeekEnd)
	}

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			var nmsg *logspray.Message
			var err error
			if !reverse {
				nmsg, err = readMessageFromFile(r)
			} else {
				nmsg, err = readMessageFromFileEnd(r)
			}

			if err == io.EOF {
				return nil
			}
			if err != nil {
				return errors.Wrapf(err, "failed to read message from %s", s.fn)
			}

			t, _ := ptypes.Timestamp(nmsg.Time)
			if t.Before(from) || t.After(to) {
				continue
			}
			if !matcher(hdr, nmsg) {
				continue
			}
			if offset > 0 {
				offset--
				continue
			}
			msgFunc(nmsg)
			count--
			if count == 0 {
				return nil
			}
		}
	}
}

func (s *ShardFile) writeMessageToFile(ctx context.Context, m *logspray.Message) error {
	s.Lock()
	defer s.Unlock()
	var err error

	if s.writer == nil {
		dir := filepath.Dir(s.fn)
		err := os.MkdirAll(dir, 0777)
		s.writer, err = os.Create(s.fn)
		if err != nil {
			return err
		}
		if s.pb {
			hm := &logspray.Message{
				ControlMessage: logspray.Message_SETHEADER,
				StreamID:       m.StreamID,
				Time:           m.Time,
				Labels:         s.labels,
			}
			sz, err := writePBMessageToFile(s.writer, hm)
			if err != nil {
				s.writer = nil
				return errors.Wrapf(err, "failed writing file header")
			}
			s.offset += int64(sz)
		}
	}

	if !s.pb {
		mt, _ := ptypes.Timestamp(m.Time)
		_, err = fmt.Fprintf(s.writer, "%s %s\n", mt.Format(time.RFC3339Nano), m.Text)
		return errors.Wrapf(err, "failed writing raw file %s", s.fn)
	}

	sz, err := writePBMessageToFile(s.writer, m)
	if err != nil {
		return errors.Wrapf(err, "failed writing to %s", s.fn)
	}
	s.offset += int64(sz)

	return err
}

func writePBMessageToFile(w io.Writer, msg *logspray.Message) (uint32, error) {
	bs, err := proto.Marshal(msg)
	if err != nil {
		return 0, errors.Wrap(err, "marshal protobuf failed")
	}

	// We use a uin16 for the size field, we'll drop anything that's
	// too big here.
	if len(bs) > 65535 {
		return 0, errors.Errorf("marshaled protobuf is %d bytes too large", len(bs)-65535)
	}

	buf := bytes.NewBuffer([]byte{})
	binary.Write(buf, binary.LittleEndian, uint16(len(bs)))
	buf.Write(bs)
	binary.Write(buf, binary.LittleEndian, uint16(len(bs)))

	n, err := w.Write(buf.Bytes())
	if n != len(buf.Bytes()) || err != nil {
		return 0, errors.Wrap(err, "write pb to file failed")
	}

	return uint32(len(buf.Bytes())), nil
}

func readMessageFromFile(r io.Reader) (*logspray.Message, error) {
	szbs := make([]byte, 2)
	szbuf := bytes.NewBuffer(szbs)
	n, err := r.Read(szbs)
	if err == io.EOF {
		return nil, err
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed read message size header")
	}
	if n != len(szbs) {
		return nil, errors.New("short read for message size header")
	}

	szbytes := uint16(0)
	err = binary.Read(szbuf, binary.LittleEndian, &szbytes)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to decode message size")
	}

	// We read the protobuf, and the trailing size bytes.
	pbbs := make([]byte, szbytes+2)
	n, err = r.Read(pbbs)
	if err == io.EOF {
		return nil, err
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed read message")
	}
	if n != len(pbbs) {
		return nil, errors.New("short read for message")
	}

	trailSzbytes := uint16(0)
	szbuf = bytes.NewBuffer(pbbs[len(pbbs)-2:])
	err = binary.Read(szbuf, binary.LittleEndian, &trailSzbytes)
	if err != nil {
		return nil, errors.Wrap(err, "failed reading trail size")
	}

	if szbytes != trailSzbytes {
		return nil, errors.Wrap(err, "header and trail size mismatch")
	}

	msg := logspray.Message{}
	err = proto.Unmarshal(pbbs[:len(pbbs)-2], &msg)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to umarshal log message proto")
	}

	return &msg, nil
}

// readMessageFromFileEnd tries to read a message before the current position
// of the io.ReadSeeker
func readMessageFromFileEnd(r io.ReadSeeker) (*logspray.Message, error) {
	if off, _ := r.Seek(0, io.SeekCurrent); off == 0 {
		return nil, io.EOF
	}

	_, err := r.Seek(-2, io.SeekCurrent)
	if err != nil {
		return nil, errors.Wrapf(err, "could not seek back to message trailer")
	}

	szbs := make([]byte, 2)
	szbuf := bytes.NewBuffer(szbs)
	n, err := r.Read(szbs)
	if err == io.EOF {
		return nil, err
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed read message size trailer")
	}
	if n != len(szbs) {
		return nil, errors.New("short read for message size trailer")
	}

	szbytes := uint16(0)
	err = binary.Read(szbuf, binary.LittleEndian, &szbytes)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to decode message size trailer")
	}

	start, err := r.Seek(-1*int64(4+szbytes), io.SeekCurrent)
	if err != nil {
		return nil, errors.Wrapf(err, "could not seek back to message header")
	}

	msg, err := readMessageFromFile(r)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to read message")
	}

	_, err = r.Seek(start, io.SeekStart)

	return msg, nil
}
<commit_msg>shardfile: fix crash on nil timestamp<commit_after>// Copyright 2016 Qubit Digital Ltd.
// 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 logspray is a collection of tools for streaming and indexing
// large volumes of dynamic logs.

package indexer

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/QubitProducts/logspray/proto/logspray"
	"github.com/QubitProducts/logspray/ql"
	"github.com/gogo/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"github.com/pkg/errors"
)

// Search searches the shard file for messages in the provided time range,
// matched by matcher, and passes them to msgFunc. If the reverse is true the
// file will be searched in reverse order
func (s *ShardFile) Search(ctx context.Context, msgFunc logspray.MessageFunc, matcher ql.MatchFunc, from, to time.Time, count, offset uint64, reverse bool) error {
	s.RLock()
	defer s.RUnlock()

	r, err := os.Open(s.fn)
	defer r.Close()
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", s.fn)
	}

	hdr, err := readMessageFromFile(r)
	if err != nil {
		return errors.Wrapf(err, "failed to read file header %s", s.fn)
	}
	err = msgFunc(hdr)
	if err != nil {
		return err
	}
	if reverse {
		r.Seek(0, io.SeekEnd)
	}

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			var nmsg *logspray.Message
			var err error
			if !reverse {
				nmsg, err = readMessageFromFile(r)
			} else {
				nmsg, err = readMessageFromFileEnd(r)
			}

			if err == io.EOF {
				return nil
			}
			if err != nil {
				return errors.Wrapf(err, "failed to read message from %s", s.fn)
			}

			if nmsg.Time == nil {
				continue
			}
			t, _ := ptypes.Timestamp(nmsg.Time)
			if t.Before(from) || t.After(to) {
				continue
			}
			if !matcher(hdr, nmsg) {
				continue
			}
			if offset > 0 {
				offset--
				continue
			}
			msgFunc(nmsg)
			count--
			if count == 0 {
				return nil
			}
		}
	}
}

func (s *ShardFile) writeMessageToFile(ctx context.Context, m *logspray.Message) error {
	s.Lock()
	defer s.Unlock()
	var err error

	if s.writer == nil {
		dir := filepath.Dir(s.fn)
		err := os.MkdirAll(dir, 0777)
		s.writer, err = os.Create(s.fn)
		if err != nil {
			return err
		}
		if s.pb {
			hm := &logspray.Message{
				ControlMessage: logspray.Message_SETHEADER,
				StreamID:       m.StreamID,
				Time:           m.Time,
				Labels:         s.labels,
			}
			sz, err := writePBMessageToFile(s.writer, hm)
			if err != nil {
				s.writer = nil
				return errors.Wrapf(err, "failed writing file header")
			}
			s.offset += int64(sz)
		}
	}

	if !s.pb {
		mt, _ := ptypes.Timestamp(m.Time)
		_, err = fmt.Fprintf(s.writer, "%s %s\n", mt.Format(time.RFC3339Nano), m.Text)
		return errors.Wrapf(err, "failed writing raw file %s", s.fn)
	}

	sz, err := writePBMessageToFile(s.writer, m)
	if err != nil {
		return errors.Wrapf(err, "failed writing to %s", s.fn)
	}
	s.offset += int64(sz)

	return err
}

func writePBMessageToFile(w io.Writer, msg *logspray.Message) (uint32, error) {
	bs, err := proto.Marshal(msg)
	if err != nil {
		return 0, errors.Wrap(err, "marshal protobuf failed")
	}

	// We use a uin16 for the size field, we'll drop anything that's
	// too big here.
	if len(bs) > 65535 {
		return 0, errors.Errorf("marshaled protobuf is %d bytes too large", len(bs)-65535)
	}

	buf := bytes.NewBuffer([]byte{})
	binary.Write(buf, binary.LittleEndian, uint16(len(bs)))
	buf.Write(bs)
	binary.Write(buf, binary.LittleEndian, uint16(len(bs)))

	n, err := w.Write(buf.Bytes())
	if n != len(buf.Bytes()) || err != nil {
		return 0, errors.Wrap(err, "write pb to file failed")
	}

	return uint32(len(buf.Bytes())), nil
}

func readMessageFromFile(r io.Reader) (*logspray.Message, error) {
	szbs := make([]byte, 2)
	szbuf := bytes.NewBuffer(szbs)
	n, err := r.Read(szbs)
	if err == io.EOF {
		return nil, err
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed read message size header")
	}
	if n != len(szbs) {
		return nil, errors.New("short read for message size header")
	}

	szbytes := uint16(0)
	err = binary.Read(szbuf, binary.LittleEndian, &szbytes)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to decode message size")
	}

	// We read the protobuf, and the trailing size bytes.
	pbbs := make([]byte, szbytes+2)
	n, err = r.Read(pbbs)
	if err == io.EOF {
		return nil, err
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed read message")
	}
	if n != len(pbbs) {
		return nil, errors.New("short read for message")
	}

	trailSzbytes := uint16(0)
	szbuf = bytes.NewBuffer(pbbs[len(pbbs)-2:])
	err = binary.Read(szbuf, binary.LittleEndian, &trailSzbytes)
	if err != nil {
		return nil, errors.Wrap(err, "failed reading trail size")
	}

	if szbytes != trailSzbytes {
		return nil, errors.Wrap(err, "header and trail size mismatch")
	}

	msg := logspray.Message{}
	err = proto.Unmarshal(pbbs[:len(pbbs)-2], &msg)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to umarshal log message proto")
	}

	return &msg, nil
}

// readMessageFromFileEnd tries to read a message before the current position
// of the io.ReadSeeker
func readMessageFromFileEnd(r io.ReadSeeker) (*logspray.Message, error) {
	if off, _ := r.Seek(0, io.SeekCurrent); off == 0 {
		return nil, io.EOF
	}

	_, err := r.Seek(-2, io.SeekCurrent)
	if err != nil {
		return nil, errors.Wrapf(err, "could not seek back to message trailer")
	}

	szbs := make([]byte, 2)
	szbuf := bytes.NewBuffer(szbs)
	n, err := r.Read(szbs)
	if err == io.EOF {
		return nil, err
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed read message size trailer")
	}
	if n != len(szbs) {
		return nil, errors.New("short read for message size trailer")
	}

	szbytes := uint16(0)
	err = binary.Read(szbuf, binary.LittleEndian, &szbytes)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to decode message size trailer")
	}

	start, err := r.Seek(-1*int64(4+szbytes), io.SeekCurrent)
	if err != nil {
		return nil, errors.Wrapf(err, "could not seek back to message header")
	}

	msg, err := readMessageFromFile(r)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to read message")
	}

	_, err = r.Seek(start, io.SeekStart)

	return msg, nil
}
<|endoftext|>


<commit_before>package pq

import (
	"database/sql/driver"
	"encoding/binary"
	"errors"
	"fmt"
	"sync"
)

var (
	errCopyInClosed               = errors.New("pq: copyin statement has already been closed")
	errBinaryCopyNotSupported     = errors.New("pq: only text format supported for COPY")
	errCopyToNotSupported         = errors.New("pq: COPY TO is not supported")
	errCopyNotSupportedOutsideTxn = errors.New("pq: COPY is only allowed inside a transaction")
	errCopyInProgress             = errors.New("pq: COPY in progress")
)

// CopyIn creates a COPY FROM statement which can be prepared with
// Tx.Prepare().  The target table should be visible in search_path.
func CopyIn(table string, columns ...string) string {
	stmt := "COPY " + QuoteIdentifier(table) + " ("
	for i, col := range columns {
		if i != 0 {
			stmt += ", "
		}
		stmt += QuoteIdentifier(col)
	}
	stmt += ") FROM STDIN"
	return stmt
}

// CopyInSchema creates a COPY FROM statement which can be prepared with
// Tx.Prepare().
func CopyInSchema(schema, table string, columns ...string) string {
	stmt := "COPY " + QuoteIdentifier(schema) + "." + QuoteIdentifier(table) + " ("
	for i, col := range columns {
		if i != 0 {
			stmt += ", "
		}
		stmt += QuoteIdentifier(col)
	}
	stmt += ") FROM STDIN"
	return stmt
}

type copyin struct {
	cn      *conn
	buffer  []byte
	rowData chan []byte
	done    chan bool

	closed bool

	mu struct {
		sync.Mutex
		err error
		driver.Result
	}
}

const ciBufferSize = 64 * 1024

// flush buffer before the buffer is filled up and needs reallocation
const ciBufferFlushSize = 63 * 1024

func (cn *conn) prepareCopyIn(q string) (_ driver.Stmt, err error) {
	if !cn.isInTransaction() {
		return nil, errCopyNotSupportedOutsideTxn
	}

	ci := &copyin{
		cn:      cn,
		buffer:  make([]byte, 0, ciBufferSize),
		rowData: make(chan []byte),
		done:    make(chan bool, 1),
	}
	// add CopyData identifier + 4 bytes for message length
	ci.buffer = append(ci.buffer, 'd', 0, 0, 0, 0)

	b := cn.writeBuf('Q')
	b.string(q)
	cn.send(b)

awaitCopyInResponse:
	for {
		t, r := cn.recv1()
		switch t {
		case 'G':
			if r.byte() != 0 {
				err = errBinaryCopyNotSupported
				break awaitCopyInResponse
			}
			go ci.resploop()
			return ci, nil
		case 'H':
			err = errCopyToNotSupported
			break awaitCopyInResponse
		case 'E':
			err = parseError(r)
		case 'Z':
			if err == nil {
				ci.setBad(driver.ErrBadConn)
				errorf("unexpected ReadyForQuery in response to COPY")
			}
			cn.processReadyForQuery(r)
			return nil, err
		default:
			ci.setBad(driver.ErrBadConn)
			errorf("unknown response for copy query: %q", t)
		}
	}

	// something went wrong, abort COPY before we return
	b = cn.writeBuf('f')
	b.string(err.Error())
	cn.send(b)

	for {
		t, r := cn.recv1()
		switch t {
		case 'c', 'C', 'E':
		case 'Z':
			// correctly aborted, we're done
			cn.processReadyForQuery(r)
			return nil, err
		default:
			ci.setBad(driver.ErrBadConn)
			errorf("unknown response for CopyFail: %q", t)
		}
	}
}

func (ci *copyin) flush(buf []byte) {
	// set message length (without message identifier)
	binary.BigEndian.PutUint32(buf[1:], uint32(len(buf)-1))

	_, err := ci.cn.c.Write(buf)
	if err != nil {
		panic(err)
	}
}

func (ci *copyin) resploop() {
	for {
		var r readBuf
		t, err := ci.cn.recvMessage(&r)
		if err != nil {
			ci.setBad(driver.ErrBadConn)
			ci.setError(err)
			ci.done <- true
			return
		}
		switch t {
		case 'C':
			// complete
			res, _ := ci.cn.parseComplete(r.string())
			ci.setResult(res)
		case 'N':
			if n := ci.cn.noticeHandler; n != nil {
				n(parseError(&r))
			}
		case 'Z':
			ci.cn.processReadyForQuery(&r)
			ci.done <- true
			return
		case 'E':
			err := parseError(&r)
			ci.setError(err)
		default:
			ci.setBad(driver.ErrBadConn)
			ci.setError(fmt.Errorf("unknown response during CopyIn: %q", t))
			ci.done <- true
			return
		}
	}
}

func (ci *copyin) setBad(err error) {
	ci.cn.err.set(err)
}

func (ci *copyin) getBad() error {
	return ci.cn.err.get()
}

func (ci *copyin) err() error {
	ci.mu.Lock()
	err := ci.mu.err
	ci.mu.Unlock()
	return err
}

// setError() sets ci.err if one has not been set already.  Caller must not be
// holding ci.Mutex.
func (ci *copyin) setError(err error) {
	ci.mu.Lock()
	if ci.mu.err == nil {
		ci.mu.err = err
	}
	ci.mu.Unlock()
}

func (ci *copyin) setResult(result driver.Result) {
	ci.mu.Lock()
	ci.mu.Result = result
	ci.mu.Unlock()
}

func (ci *copyin) getResult() driver.Result {
	ci.mu.Lock()
	result := ci.mu.Result
	ci.mu.Unlock()
	if result == nil {
		return driver.RowsAffected(0)
	}
	return result
}

func (ci *copyin) NumInput() int {
	return -1
}

func (ci *copyin) Query(v []driver.Value) (r driver.Rows, err error) {
	return nil, ErrNotSupported
}

// Exec inserts values into the COPY stream. The insert is asynchronous
// and Exec can return errors from previous Exec calls to the same
// COPY stmt.
//
// You need to call Exec(nil) to sync the COPY stream and to get any
// errors from pending data, since Stmt.Close() doesn't return errors
// to the user.
func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
	if ci.closed {
		return nil, errCopyInClosed
	}

	if err := ci.getBad(); err != nil {
		return nil, err
	}
	defer ci.cn.errRecover(&err)

	if err := ci.err(); err != nil {
		return nil, err
	}

	if len(v) == 0 {
		if err := ci.Close(); err != nil {
			return driver.RowsAffected(0), err
		}

		return ci.getResult(), nil
	}

	numValues := len(v)
	for i, value := range v {
		ci.buffer = appendEncodedText(&ci.cn.parameterStatus, ci.buffer, value)
		if i < numValues-1 {
			ci.buffer = append(ci.buffer, '\t')
		}
	}

	ci.buffer = append(ci.buffer, '\n')

	if len(ci.buffer) > ciBufferFlushSize {
		ci.flush(ci.buffer)
		// reset buffer, keep bytes for message identifier and length
		ci.buffer = ci.buffer[:5]
	}

	return driver.RowsAffected(0), nil
}

func (ci *copyin) Close() (err error) {
	if ci.closed { // Don't do anything, we're already closed
		return nil
	}
	ci.closed = true

	if err := ci.getBad(); err != nil {
		return err
	}
	defer ci.cn.errRecover(&err)

	if len(ci.buffer) > 0 {
		ci.flush(ci.buffer)
	}
	// Avoid touching the scratch buffer as resploop could be using it.
	err = ci.cn.sendSimpleMessage('c')
	if err != nil {
		return err
	}

	<-ci.done
	ci.cn.inCopy = false

	if err := ci.err(); err != nil {
		return err
	}
	return nil
}
<commit_msg>expose raw CopyData command (#1077)<commit_after>package pq

import (
	"context"
	"database/sql/driver"
	"encoding/binary"
	"errors"
	"fmt"
	"sync"
)

var (
	errCopyInClosed               = errors.New("pq: copyin statement has already been closed")
	errBinaryCopyNotSupported     = errors.New("pq: only text format supported for COPY")
	errCopyToNotSupported         = errors.New("pq: COPY TO is not supported")
	errCopyNotSupportedOutsideTxn = errors.New("pq: COPY is only allowed inside a transaction")
	errCopyInProgress             = errors.New("pq: COPY in progress")
)

// CopyIn creates a COPY FROM statement which can be prepared with
// Tx.Prepare().  The target table should be visible in search_path.
func CopyIn(table string, columns ...string) string {
	stmt := "COPY " + QuoteIdentifier(table) + " ("
	for i, col := range columns {
		if i != 0 {
			stmt += ", "
		}
		stmt += QuoteIdentifier(col)
	}
	stmt += ") FROM STDIN"
	return stmt
}

// CopyInSchema creates a COPY FROM statement which can be prepared with
// Tx.Prepare().
func CopyInSchema(schema, table string, columns ...string) string {
	stmt := "COPY " + QuoteIdentifier(schema) + "." + QuoteIdentifier(table) + " ("
	for i, col := range columns {
		if i != 0 {
			stmt += ", "
		}
		stmt += QuoteIdentifier(col)
	}
	stmt += ") FROM STDIN"
	return stmt
}

type copyin struct {
	cn      *conn
	buffer  []byte
	rowData chan []byte
	done    chan bool

	closed bool

	mu struct {
		sync.Mutex
		err error
		driver.Result
	}
}

const ciBufferSize = 64 * 1024

// flush buffer before the buffer is filled up and needs reallocation
const ciBufferFlushSize = 63 * 1024

func (cn *conn) prepareCopyIn(q string) (_ driver.Stmt, err error) {
	if !cn.isInTransaction() {
		return nil, errCopyNotSupportedOutsideTxn
	}

	ci := &copyin{
		cn:      cn,
		buffer:  make([]byte, 0, ciBufferSize),
		rowData: make(chan []byte),
		done:    make(chan bool, 1),
	}
	// add CopyData identifier + 4 bytes for message length
	ci.buffer = append(ci.buffer, 'd', 0, 0, 0, 0)

	b := cn.writeBuf('Q')
	b.string(q)
	cn.send(b)

awaitCopyInResponse:
	for {
		t, r := cn.recv1()
		switch t {
		case 'G':
			if r.byte() != 0 {
				err = errBinaryCopyNotSupported
				break awaitCopyInResponse
			}
			go ci.resploop()
			return ci, nil
		case 'H':
			err = errCopyToNotSupported
			break awaitCopyInResponse
		case 'E':
			err = parseError(r)
		case 'Z':
			if err == nil {
				ci.setBad(driver.ErrBadConn)
				errorf("unexpected ReadyForQuery in response to COPY")
			}
			cn.processReadyForQuery(r)
			return nil, err
		default:
			ci.setBad(driver.ErrBadConn)
			errorf("unknown response for copy query: %q", t)
		}
	}

	// something went wrong, abort COPY before we return
	b = cn.writeBuf('f')
	b.string(err.Error())
	cn.send(b)

	for {
		t, r := cn.recv1()
		switch t {
		case 'c', 'C', 'E':
		case 'Z':
			// correctly aborted, we're done
			cn.processReadyForQuery(r)
			return nil, err
		default:
			ci.setBad(driver.ErrBadConn)
			errorf("unknown response for CopyFail: %q", t)
		}
	}
}

func (ci *copyin) flush(buf []byte) {
	// set message length (without message identifier)
	binary.BigEndian.PutUint32(buf[1:], uint32(len(buf)-1))

	_, err := ci.cn.c.Write(buf)
	if err != nil {
		panic(err)
	}
}

func (ci *copyin) resploop() {
	for {
		var r readBuf
		t, err := ci.cn.recvMessage(&r)
		if err != nil {
			ci.setBad(driver.ErrBadConn)
			ci.setError(err)
			ci.done <- true
			return
		}
		switch t {
		case 'C':
			// complete
			res, _ := ci.cn.parseComplete(r.string())
			ci.setResult(res)
		case 'N':
			if n := ci.cn.noticeHandler; n != nil {
				n(parseError(&r))
			}
		case 'Z':
			ci.cn.processReadyForQuery(&r)
			ci.done <- true
			return
		case 'E':
			err := parseError(&r)
			ci.setError(err)
		default:
			ci.setBad(driver.ErrBadConn)
			ci.setError(fmt.Errorf("unknown response during CopyIn: %q", t))
			ci.done <- true
			return
		}
	}
}

func (ci *copyin) setBad(err error) {
	ci.cn.err.set(err)
}

func (ci *copyin) getBad() error {
	return ci.cn.err.get()
}

func (ci *copyin) err() error {
	ci.mu.Lock()
	err := ci.mu.err
	ci.mu.Unlock()
	return err
}

// setError() sets ci.err if one has not been set already.  Caller must not be
// holding ci.Mutex.
func (ci *copyin) setError(err error) {
	ci.mu.Lock()
	if ci.mu.err == nil {
		ci.mu.err = err
	}
	ci.mu.Unlock()
}

func (ci *copyin) setResult(result driver.Result) {
	ci.mu.Lock()
	ci.mu.Result = result
	ci.mu.Unlock()
}

func (ci *copyin) getResult() driver.Result {
	ci.mu.Lock()
	result := ci.mu.Result
	ci.mu.Unlock()
	if result == nil {
		return driver.RowsAffected(0)
	}
	return result
}

func (ci *copyin) NumInput() int {
	return -1
}

func (ci *copyin) Query(v []driver.Value) (r driver.Rows, err error) {
	return nil, ErrNotSupported
}

// Exec inserts values into the COPY stream. The insert is asynchronous
// and Exec can return errors from previous Exec calls to the same
// COPY stmt.
//
// You need to call Exec(nil) to sync the COPY stream and to get any
// errors from pending data, since Stmt.Close() doesn't return errors
// to the user.
func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
	if ci.closed {
		return nil, errCopyInClosed
	}

	if err := ci.getBad(); err != nil {
		return nil, err
	}
	defer ci.cn.errRecover(&err)

	if err := ci.err(); err != nil {
		return nil, err
	}

	if len(v) == 0 {
		if err := ci.Close(); err != nil {
			return driver.RowsAffected(0), err
		}

		return ci.getResult(), nil
	}

	numValues := len(v)
	for i, value := range v {
		ci.buffer = appendEncodedText(&ci.cn.parameterStatus, ci.buffer, value)
		if i < numValues-1 {
			ci.buffer = append(ci.buffer, '\t')
		}
	}

	ci.buffer = append(ci.buffer, '\n')

	if len(ci.buffer) > ciBufferFlushSize {
		ci.flush(ci.buffer)
		// reset buffer, keep bytes for message identifier and length
		ci.buffer = ci.buffer[:5]
	}

	return driver.RowsAffected(0), nil
}

// CopyData executes a raw CopyData command using the PostgreSQL Frontend/Backend
// protocol. Use Exec(nil) to finish the command.
func (ci *copyin) CopyData(ctx context.Context, line string) (r driver.Result, err error) {
	if ci.closed {
		return nil, errCopyInClosed
	}

	if finish := ci.cn.watchCancel(ctx); finish != nil {
		defer finish()
	}

	if err := ci.getBad(); err != nil {
		return nil, err
	}
	defer ci.cn.errRecover(&err)

	if err := ci.err(); err != nil {
		return nil, err
	}

	ci.buffer = append(ci.buffer, []byte(line)...)
	ci.buffer = append(ci.buffer, '\n')

	if len(ci.buffer) > ciBufferFlushSize {
		ci.flush(ci.buffer)
		// reset buffer, keep bytes for message identifier and length
		ci.buffer = ci.buffer[:5]
	}

	return driver.RowsAffected(0), nil
}

func (ci *copyin) Close() (err error) {
	if ci.closed { // Don't do anything, we're already closed
		return nil
	}
	ci.closed = true

	if err := ci.getBad(); err != nil {
		return err
	}
	defer ci.cn.errRecover(&err)

	if len(ci.buffer) > 0 {
		ci.flush(ci.buffer)
	}
	// Avoid touching the scratch buffer as resploop could be using it.
	err = ci.cn.sendSimpleMessage('c')
	if err != nil {
		return err
	}

	<-ci.done
	ci.cn.inCopy = false

	if err := ci.err(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package zapdriver

import (
	"strings"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// Core is a zapdriver specific core wrapped around the default zap core. It
// allows to merge all defined labels
type core struct {
	zapcore.Core

	// permLabels is a collection of labels that have been added to the logger
	// through the use of `With()`. These labels should never be cleared after
	// logging a single entry, unlike `tempLabel`.
	permLabels *labels

	// tempLabels keeps a record of all the labels that need to be applied to the
	// current log entry. Zap serializes log fields at different parts of the
	// stack, one such location is when calling `core.With` and the other one is
	// when calling `core.Write`. This makes it impossible to (for example) take
	// all `labels.xxx` fields, and wrap them in the `labels` namespace in one go.
	//
	// Instead, we have to filter out these labels at both locations, and then add
	// them back in the proper format right before we call `Write` on the original
	// Zap core.
	tempLabels *labels
}

// WrapCore returns a `zap.Option` that wraps the default core with the
// zapdriver one.
func WrapCore() zap.Option {
	return zap.WrapCore(func(c zapcore.Core) zapcore.Core {
		return &core{c, newLabels(), newLabels()}
	})
}

// With adds structured context to the Core.
func (c *core) With(fields []zap.Field) zapcore.Core {
	var lbls *labels
	lbls, fields = c.extractLabels(fields)

	lbls.mutex.Lock()
	for k, v := range lbls.store {
		c.permLabels.store[k] = v
	}
	lbls.mutex.Unlock()

	return &core{c.Core.With(fields), c.permLabels, newLabels()}
}

// Check determines whether the supplied Entry should be logged (using the
// embedded LevelEnabler and possibly some extra logic). If the entry
// should be logged, the Core adds itself to the CheckedEntry and returns
// the result.
//
// Callers must use Check before calling Write.
func (c *core) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
	if c.Enabled(ent.Level) {
		return ce.AddCore(ent, c)
	}

	return ce
}

func (c *core) Write(ent zapcore.Entry, fields []zapcore.Field) error {
	var lbls *labels
	lbls, fields = c.extractLabels(fields)

	lbls.mutex.Lock()
	for k, v := range lbls.store {
		c.tempLabels.store[k] = v
	}
	lbls.mutex.Unlock()

	fields = append(fields, labelsField(c.allLabels()))
	fields = c.withSourceLocation(ent, fields)

	c.tempLabels = newLabels()

	return c.Core.Write(ent, fields)
}

// Sync flushes buffered logs (if any).
func (c *core) Sync() error {
	return c.Core.Sync()
}

func (c *core) allLabels() *labels {
	lbls := newLabels()

	lbls.mutex.Lock()
	for k, v := range c.permLabels.store {
		lbls.store[k] = v
	}

	for k, v := range c.tempLabels.store {
		lbls.store[k] = v
	}
	lbls.mutex.Unlock()

	return lbls
}

func (c *core) extractLabels(fields []zapcore.Field) (*labels, []zapcore.Field) {
	lbls := newLabels()
	out := []zapcore.Field{}

	lbls.mutex.Lock()
	for i := range fields {
		if !isLabelField(fields[i]) {
			out = append(out, fields[i])
			continue
		}

		lbls.store[strings.Replace(fields[i].Key, "labels.", "", 1)] = fields[i].String
	}
	lbls.mutex.Unlock()

	return lbls, out
}

func (c *core) withLabels(fields []zapcore.Field) []zapcore.Field {
	lbls := newLabels()
	out := []zapcore.Field{}

	lbls.mutex.Lock()
	for i := range fields {
		if isLabelField(fields[i]) {
			lbls.store[strings.Replace(fields[i].Key, "labels.", "", 1)] = fields[i].String
			continue
		}

		out = append(out, fields[i])
	}
	lbls.mutex.Unlock()

	return append(out, labelsField(lbls))
}

func (c *core) withSourceLocation(ent zapcore.Entry, fields []zapcore.Field) []zapcore.Field {
	// If the source location was manually set, don't overwrite it
	for i := range fields {
		if fields[i].Key == sourceKey {
			return fields
		}
	}

	if !ent.Caller.Defined {
		return fields
	}

	return append(fields, SourceLocation(ent.Caller.PC, ent.Caller.File, ent.Caller.Line, true))
}
<commit_msg>add missing mutex locks<commit_after>package zapdriver

import (
	"strings"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

// Core is a zapdriver specific core wrapped around the default zap core. It
// allows to merge all defined labels
type core struct {
	zapcore.Core

	// permLabels is a collection of labels that have been added to the logger
	// through the use of `With()`. These labels should never be cleared after
	// logging a single entry, unlike `tempLabel`.
	permLabels *labels

	// tempLabels keeps a record of all the labels that need to be applied to the
	// current log entry. Zap serializes log fields at different parts of the
	// stack, one such location is when calling `core.With` and the other one is
	// when calling `core.Write`. This makes it impossible to (for example) take
	// all `labels.xxx` fields, and wrap them in the `labels` namespace in one go.
	//
	// Instead, we have to filter out these labels at both locations, and then add
	// them back in the proper format right before we call `Write` on the original
	// Zap core.
	tempLabels *labels
}

// WrapCore returns a `zap.Option` that wraps the default core with the
// zapdriver one.
func WrapCore() zap.Option {
	return zap.WrapCore(func(c zapcore.Core) zapcore.Core {
		return &core{c, newLabels(), newLabels()}
	})
}

// With adds structured context to the Core.
func (c *core) With(fields []zap.Field) zapcore.Core {
	var lbls *labels
	lbls, fields = c.extractLabels(fields)

	lbls.mutex.RLock()
	c.permLabels.mutex.Lock()
	for k, v := range lbls.store {
		c.permLabels.store[k] = v
	}
	c.permLabels.mutex.Unlock()
	lbls.mutex.RUnlock()

	return &core{c.Core.With(fields), c.permLabels, newLabels()}
}

// Check determines whether the supplied Entry should be logged (using the
// embedded LevelEnabler and possibly some extra logic). If the entry
// should be logged, the Core adds itself to the CheckedEntry and returns
// the result.
//
// Callers must use Check before calling Write.
func (c *core) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
	if c.Enabled(ent.Level) {
		return ce.AddCore(ent, c)
	}

	return ce
}

func (c *core) Write(ent zapcore.Entry, fields []zapcore.Field) error {
	var lbls *labels
	lbls, fields = c.extractLabels(fields)

	lbls.mutex.RLock()
	c.tempLabels.mutex.Lock()
	for k, v := range lbls.store {
		c.tempLabels.store[k] = v
	}
	c.tempLabels.mutex.Unlock()
	lbls.mutex.RUnlock()

	fields = append(fields, labelsField(c.allLabels()))
	fields = c.withSourceLocation(ent, fields)

	c.tempLabels = newLabels()

	return c.Core.Write(ent, fields)
}

// Sync flushes buffered logs (if any).
func (c *core) Sync() error {
	return c.Core.Sync()
}

func (c *core) allLabels() *labels {
	lbls := newLabels()

	lbls.mutex.Lock()
	c.permLabels.mutex.RLock()
	for k, v := range c.permLabels.store {
		lbls.store[k] = v
	}
	c.permLabels.mutex.RUnlock()

	c.tempLabels.mutex.RLock()
	for k, v := range c.tempLabels.store {
		lbls.store[k] = v
	}
	c.tempLabels.mutex.RUnlock()
	lbls.mutex.Unlock()

	return lbls
}

func (c *core) extractLabels(fields []zapcore.Field) (*labels, []zapcore.Field) {
	lbls := newLabels()
	out := []zapcore.Field{}

	lbls.mutex.Lock()
	for i := range fields {
		if !isLabelField(fields[i]) {
			out = append(out, fields[i])
			continue
		}

		lbls.store[strings.Replace(fields[i].Key, "labels.", "", 1)] = fields[i].String
	}
	lbls.mutex.Unlock()

	return lbls, out
}

func (c *core) withLabels(fields []zapcore.Field) []zapcore.Field {
	lbls := newLabels()
	out := []zapcore.Field{}

	lbls.mutex.Lock()
	for i := range fields {
		if isLabelField(fields[i]) {
			lbls.store[strings.Replace(fields[i].Key, "labels.", "", 1)] = fields[i].String
			continue
		}

		out = append(out, fields[i])
	}
	lbls.mutex.Unlock()

	return append(out, labelsField(lbls))
}

func (c *core) withSourceLocation(ent zapcore.Entry, fields []zapcore.Field) []zapcore.Field {
	// If the source location was manually set, don't overwrite it
	for i := range fields {
		if fields[i].Key == sourceKey {
			return fields
		}
	}

	if !ent.Caller.Defined {
		return fields
	}

	return append(fields, SourceLocation(ent.Caller.PC, ent.Caller.File, ent.Caller.Line, true))
}
<|endoftext|>