<commit_before>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcwire_test

import (
	"bytes"
	"github.com/conformal/btcwire"
	"github.com/davecgh/go-spew/spew"
	"reflect"
	"testing"
)

// TestGetHeaders tests the MsgGetHeader API.
func TestGetHeaders(t *testing.T) {
	pver := btcwire.ProtocolVersion

	// Block 99500 hash.
	hashStr := "000000000002e7ad7b9eef9479e4aabc65cb831269cc20d2632c13684406dee0"
	locatorHash, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Ensure the command is expected value.
	wantCmd := "getheaders"
	msg := btcwire.NewMsgGetHeaders()
	if cmd := msg.Command(); cmd != wantCmd {
		t.Errorf("NewMsgGetHeaders: wrong command - got %v want %v",
			cmd, wantCmd)
	}

	// Ensure max payload is expected value for latest protocol version.
	// Protocol version 4 bytes + num hashes (varInt) + max block locator
	// hashes + hash stop.
	wantPayload := uint32(16045)
	maxPayload := msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}

	// Ensure block locator hashes are added properly.
	err = msg.AddBlockLocatorHash(locatorHash)
	if err != nil {
		t.Errorf("AddBlockLocatorHash: %v", err)
	}
	if msg.BlockLocatorHashes[0] != locatorHash {
		t.Errorf("AddBlockLocatorHash: wrong block locator added - "+
			"got %v, want %v",
			spew.Sprint(msg.BlockLocatorHashes[0]),
			spew.Sprint(locatorHash))
	}

	// Ensure adding more than the max allowed block locator hashes per
	// message returns an error.
	for i := 0; i < btcwire.MaxBlockLocatorsPerMsg; i++ {
		err = msg.AddBlockLocatorHash(locatorHash)
	}
	if err == nil {
		t.Errorf("AddBlockLocatorHash: expected error on too many " +
			"block locator hashes not received")
	}

	return
}

// TestGetHeadersWire tests the MsgGetHeaders wire encode and decode for various
// numbers of block locator hashes and protocol versions.
func TestGetHeadersWire(t *testing.T) {
	// Set protocol inside getheaders message.
	pver := uint32(60002)

	// Block 99499 hash.
	hashStr := "2710f40c87ec93d010a6fd95f42c59a2cbacc60b18cf6b7957535"
	hashLocator, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Block 99500 hash.
	hashStr = "2e7ad7b9eef9479e4aabc65cb831269cc20d2632c13684406dee0"
	hashLocator2, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Block 100000 hash.
	hashStr = "3ba27aa200b1cecaad478d2b00432346c3f1f3986da1afd33e506"
	hashStop, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// MsgGetHeaders message with no block locators or stop hash.
	NoLocators := btcwire.NewMsgGetHeaders()
	NoLocators.ProtocolVersion = pver
	NoLocatorsEncoded := []byte{
		0x62, 0xea, 0x00, 0x00, // Protocol version 60002
		0x00, // Varint for number of block locator hashes
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Hash stop
	}

	// MsgGetHeaders message with multiple block locators and a stop hash.
	MultiLocators := btcwire.NewMsgGetHeaders()
	MultiLocators.ProtocolVersion = pver
	MultiLocators.HashStop = *hashStop
	MultiLocators.AddBlockLocatorHash(hashLocator2)
	MultiLocators.AddBlockLocatorHash(hashLocator)
	MultiLocatorsEncoded := []byte{
		0x62, 0xea, 0x00, 0x00, // Protocol version 60002
		0x02, // Varint for number of block locator hashes
		0xe0, 0xde, 0x06, 0x44, 0x68, 0x13, 0x2c, 0x63,
		0xd2, 0x20, 0xcc, 0x69, 0x12, 0x83, 0xcb, 0x65,
		0xbc, 0xaa, 0xe4, 0x79, 0x94, 0xef, 0x9e, 0x7b,
		0xad, 0xe7, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // Block 99500 hash
		0x35, 0x75, 0x95, 0xb7, 0xf6, 0x8c, 0xb1, 0x60,
		0xcc, 0xba, 0x2c, 0x9a, 0xc5, 0x42, 0x5f, 0xd9,
		0x6f, 0x0a, 0x01, 0x3d, 0xc9, 0x7e, 0xc8, 0x40,
		0x0f, 0x71, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // Block 99499 hash
		0x06, 0xe5, 0x33, 0xfd, 0x1a, 0xda, 0x86, 0x39,
		0x1f, 0x3f, 0x6c, 0x34, 0x32, 0x04, 0xb0, 0xd2,
		0x78, 0xd4, 0xaa, 0xec, 0x1c, 0x0b, 0x20, 0xaa,
		0x27, 0xba, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, // Hash stop
	}

	tests := []struct {
		in   *btcwire.MsgGetHeaders // Message to encode
		out  *btcwire.MsgGetHeaders // Expected decoded message
		buf  []byte                 // Wire encoding
		pver uint32                 // Protocol version for wire encoding
	}{
		// Latest protocol version with no block locators.
		{
			NoLocators,
			NoLocators,
			NoLocatorsEncoded,
			btcwire.ProtocolVersion,
		},

		// Latest protocol version with multiple block locators.
		{
			MultiLocators,
			MultiLocators,
			MultiLocatorsEncoded,
			btcwire.ProtocolVersion,
		},

		// Protocol version BIP0035Version with no block locators.
		{
			NoLocators,
			NoLocators,
			NoLocatorsEncoded,
			btcwire.BIP0035Version,
		},

		// Protocol version BIP0035Version with multiple block locators.
		{
			MultiLocators,
			MultiLocators,
			MultiLocatorsEncoded,
			btcwire.BIP0035Version,
		},

		// Protocol version BIP0031Version with no block locators.
		{
			NoLocators,
			NoLocators,
			NoLocatorsEncoded,
			btcwire.BIP0031Version,
		},

		// Protocol version BIP0031Versionwith multiple block locators.
		{
			MultiLocators,
			MultiLocators,
			MultiLocatorsEncoded,
			btcwire.BIP0031Version,
		},

		// Protocol version NetAddressTimeVersion with no block locators.
		{
			NoLocators,
			NoLocators,
			NoLocatorsEncoded,
			btcwire.NetAddressTimeVersion,
		},

		// Protocol version NetAddressTimeVersion multiple block locators.
		{
			MultiLocators,
			MultiLocators,
			MultiLocatorsEncoded,
			btcwire.NetAddressTimeVersion,
		},

		// Protocol version MultipleAddressVersion with no block locators.
		{
			NoLocators,
			NoLocators,
			NoLocatorsEncoded,
			btcwire.MultipleAddressVersion,
		},

		// Protocol version MultipleAddressVersion multiple block locators.
		{
			MultiLocators,
			MultiLocators,
			MultiLocatorsEncoded,
			btcwire.MultipleAddressVersion,
		},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode the message to wire format.
		var buf bytes.Buffer
		err := test.in.BtcEncode(&buf, test.pver)
		if err != nil {
			t.Errorf("BtcEncode #%d error %v", i, err)
			continue
		}
		if !bytes.Equal(buf.Bytes(), test.buf) {
			t.Errorf("BtcEncode #%d\n got: %s want: %s", i,
				spew.Sdump(buf.Bytes()), spew.Sdump(test.buf))
			continue
		}

		// Decode the message from wire format.
		var msg btcwire.MsgGetHeaders
		rbuf := bytes.NewBuffer(test.buf)
		err = msg.BtcDecode(rbuf, test.pver)
		if err != nil {
			t.Errorf("BtcDecode #%d error %v", i, err)
			continue
		}
		if !reflect.DeepEqual(&msg, test.out) {
			t.Errorf("BtcDecode #%d\n got: %s want: %s", i,
				spew.Sdump(&msg), spew.Sdump(test.out))
			continue
		}
	}
}
<commit_msg>Add negative tests for MsgGetHeaders.<commit_after>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcwire_test

import (
	"bytes"
	"github.com/conformal/btcwire"
	"github.com/davecgh/go-spew/spew"
	"io"
	"reflect"
	"testing"
)

// TestGetHeaders tests the MsgGetHeader API.
func TestGetHeaders(t *testing.T) {
	pver := btcwire.ProtocolVersion

	// Block 99500 hash.
	hashStr := "000000000002e7ad7b9eef9479e4aabc65cb831269cc20d2632c13684406dee0"
	locatorHash, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Ensure the command is expected value.
	wantCmd := "getheaders"
	msg := btcwire.NewMsgGetHeaders()
	if cmd := msg.Command(); cmd != wantCmd {
		t.Errorf("NewMsgGetHeaders: wrong command - got %v want %v",
			cmd, wantCmd)
	}

	// Ensure max payload is expected value for latest protocol version.
	// Protocol version 4 bytes + num hashes (varInt) + max block locator
	// hashes + hash stop.
	wantPayload := uint32(16045)
	maxPayload := msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}

	// Ensure block locator hashes are added properly.
	err = msg.AddBlockLocatorHash(locatorHash)
	if err != nil {
		t.Errorf("AddBlockLocatorHash: %v", err)
	}
	if msg.BlockLocatorHashes[0] != locatorHash {
		t.Errorf("AddBlockLocatorHash: wrong block locator added - "+
			"got %v, want %v",
			spew.Sprint(msg.BlockLocatorHashes[0]),
			spew.Sprint(locatorHash))
	}

	// Ensure adding more than the max allowed block locator hashes per
	// message returns an error.
	for i := 0; i < btcwire.MaxBlockLocatorsPerMsg; i++ {
		err = msg.AddBlockLocatorHash(locatorHash)
	}
	if err == nil {
		t.Errorf("AddBlockLocatorHash: expected error on too many " +
			"block locator hashes not received")
	}

	return
}

// TestGetHeadersWire tests the MsgGetHeaders wire encode and decode for various
// numbers of block locator hashes and protocol versions.
func TestGetHeadersWire(t *testing.T) {
	// Set protocol inside getheaders message.  Use protocol version 60002
	// specifically here instead of the latest because the test data is
	// using bytes encoded with that protocol version.
	pver := uint32(60002)

	// Block 99499 hash.
	hashStr := "2710f40c87ec93d010a6fd95f42c59a2cbacc60b18cf6b7957535"
	hashLocator, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Block 99500 hash.
	hashStr = "2e7ad7b9eef9479e4aabc65cb831269cc20d2632c13684406dee0"
	hashLocator2, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Block 100000 hash.
	hashStr = "3ba27aa200b1cecaad478d2b00432346c3f1f3986da1afd33e506"
	hashStop, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// MsgGetHeaders message with no block locators or stop hash.
	noLocators := btcwire.NewMsgGetHeaders()
	noLocators.ProtocolVersion = pver
	noLocatorsEncoded := []byte{
		0x62, 0xea, 0x00, 0x00, // Protocol version 60002
		0x00, // Varint for number of block locator hashes
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Hash stop
	}

	// MsgGetHeaders message with multiple block locators and a stop hash.
	multiLocators := btcwire.NewMsgGetHeaders()
	multiLocators.ProtocolVersion = pver
	multiLocators.HashStop = *hashStop
	multiLocators.AddBlockLocatorHash(hashLocator2)
	multiLocators.AddBlockLocatorHash(hashLocator)
	multiLocatorsEncoded := []byte{
		0x62, 0xea, 0x00, 0x00, // Protocol version 60002
		0x02, // Varint for number of block locator hashes
		0xe0, 0xde, 0x06, 0x44, 0x68, 0x13, 0x2c, 0x63,
		0xd2, 0x20, 0xcc, 0x69, 0x12, 0x83, 0xcb, 0x65,
		0xbc, 0xaa, 0xe4, 0x79, 0x94, 0xef, 0x9e, 0x7b,
		0xad, 0xe7, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // Block 99500 hash
		0x35, 0x75, 0x95, 0xb7, 0xf6, 0x8c, 0xb1, 0x60,
		0xcc, 0xba, 0x2c, 0x9a, 0xc5, 0x42, 0x5f, 0xd9,
		0x6f, 0x0a, 0x01, 0x3d, 0xc9, 0x7e, 0xc8, 0x40,
		0x0f, 0x71, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // Block 99499 hash
		0x06, 0xe5, 0x33, 0xfd, 0x1a, 0xda, 0x86, 0x39,
		0x1f, 0x3f, 0x6c, 0x34, 0x32, 0x04, 0xb0, 0xd2,
		0x78, 0xd4, 0xaa, 0xec, 0x1c, 0x0b, 0x20, 0xaa,
		0x27, 0xba, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, // Hash stop
	}

	tests := []struct {
		in   *btcwire.MsgGetHeaders // Message to encode
		out  *btcwire.MsgGetHeaders // Expected decoded message
		buf  []byte                 // Wire encoding
		pver uint32                 // Protocol version for wire encoding
	}{
		// Latest protocol version with no block locators.
		{
			noLocators,
			noLocators,
			noLocatorsEncoded,
			btcwire.ProtocolVersion,
		},

		// Latest protocol version with multiple block locators.
		{
			multiLocators,
			multiLocators,
			multiLocatorsEncoded,
			btcwire.ProtocolVersion,
		},

		// Protocol version BIP0035Version with no block locators.
		{
			noLocators,
			noLocators,
			noLocatorsEncoded,
			btcwire.BIP0035Version,
		},

		// Protocol version BIP0035Version with multiple block locators.
		{
			multiLocators,
			multiLocators,
			multiLocatorsEncoded,
			btcwire.BIP0035Version,
		},

		// Protocol version BIP0031Version with no block locators.
		{
			noLocators,
			noLocators,
			noLocatorsEncoded,
			btcwire.BIP0031Version,
		},

		// Protocol version BIP0031Versionwith multiple block locators.
		{
			multiLocators,
			multiLocators,
			multiLocatorsEncoded,
			btcwire.BIP0031Version,
		},

		// Protocol version NetAddressTimeVersion with no block locators.
		{
			noLocators,
			noLocators,
			noLocatorsEncoded,
			btcwire.NetAddressTimeVersion,
		},

		// Protocol version NetAddressTimeVersion multiple block locators.
		{
			multiLocators,
			multiLocators,
			multiLocatorsEncoded,
			btcwire.NetAddressTimeVersion,
		},

		// Protocol version MultipleAddressVersion with no block locators.
		{
			noLocators,
			noLocators,
			noLocatorsEncoded,
			btcwire.MultipleAddressVersion,
		},

		// Protocol version MultipleAddressVersion multiple block locators.
		{
			multiLocators,
			multiLocators,
			multiLocatorsEncoded,
			btcwire.MultipleAddressVersion,
		},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode the message to wire format.
		var buf bytes.Buffer
		err := test.in.BtcEncode(&buf, test.pver)
		if err != nil {
			t.Errorf("BtcEncode #%d error %v", i, err)
			continue
		}
		if !bytes.Equal(buf.Bytes(), test.buf) {
			t.Errorf("BtcEncode #%d\n got: %s want: %s", i,
				spew.Sdump(buf.Bytes()), spew.Sdump(test.buf))
			continue
		}

		// Decode the message from wire format.
		var msg btcwire.MsgGetHeaders
		rbuf := bytes.NewBuffer(test.buf)
		err = msg.BtcDecode(rbuf, test.pver)
		if err != nil {
			t.Errorf("BtcDecode #%d error %v", i, err)
			continue
		}
		if !reflect.DeepEqual(&msg, test.out) {
			t.Errorf("BtcDecode #%d\n got: %s want: %s", i,
				spew.Sdump(&msg), spew.Sdump(test.out))
			continue
		}
	}
}

// TestGetHeadersWireErrors performs negative tests against wire encode and
// decode of MsgGetHeaders to confirm error paths work correctly.
func TestGetHeadersWireErrors(t *testing.T) {
	// Set protocol inside getheaders message.  Use protocol version 60002
	// specifically here instead of the latest because the test data is
	// using bytes encoded with that protocol version.
	pver := uint32(60002)
	btcwireErr := &btcwire.MessageError{}

	// Block 99499 hash.
	hashStr := "2710f40c87ec93d010a6fd95f42c59a2cbacc60b18cf6b7957535"
	hashLocator, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Block 99500 hash.
	hashStr = "2e7ad7b9eef9479e4aabc65cb831269cc20d2632c13684406dee0"
	hashLocator2, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// Block 100000 hash.
	hashStr = "3ba27aa200b1cecaad478d2b00432346c3f1f3986da1afd33e506"
	hashStop, err := btcwire.NewShaHashFromStr(hashStr)
	if err != nil {
		t.Errorf("NewShaHashFromStr: %v", err)
	}

	// MsgGetHeaders message with multiple block locators and a stop hash.
	baseGetHeaders := btcwire.NewMsgGetHeaders()
	baseGetHeaders.ProtocolVersion = pver
	baseGetHeaders.HashStop = *hashStop
	baseGetHeaders.AddBlockLocatorHash(hashLocator2)
	baseGetHeaders.AddBlockLocatorHash(hashLocator)
	baseGetHeadersEncoded := []byte{
		0x62, 0xea, 0x00, 0x00, // Protocol version 60002
		0x02, // Varint for number of block locator hashes
		0xe0, 0xde, 0x06, 0x44, 0x68, 0x13, 0x2c, 0x63,
		0xd2, 0x20, 0xcc, 0x69, 0x12, 0x83, 0xcb, 0x65,
		0xbc, 0xaa, 0xe4, 0x79, 0x94, 0xef, 0x9e, 0x7b,
		0xad, 0xe7, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // Block 99500 hash
		0x35, 0x75, 0x95, 0xb7, 0xf6, 0x8c, 0xb1, 0x60,
		0xcc, 0xba, 0x2c, 0x9a, 0xc5, 0x42, 0x5f, 0xd9,
		0x6f, 0x0a, 0x01, 0x3d, 0xc9, 0x7e, 0xc8, 0x40,
		0x0f, 0x71, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, // Block 99499 hash
		0x06, 0xe5, 0x33, 0xfd, 0x1a, 0xda, 0x86, 0x39,
		0x1f, 0x3f, 0x6c, 0x34, 0x32, 0x04, 0xb0, 0xd2,
		0x78, 0xd4, 0xaa, 0xec, 0x1c, 0x0b, 0x20, 0xaa,
		0x27, 0xba, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, // Hash stop
	}

	// Message that forces an error by having more than the max allowed
	// block locator hashes.
	maxGetHeaders := btcwire.NewMsgGetHeaders()
	for i := 0; i < btcwire.MaxBlockLocatorsPerMsg; i++ {
		maxGetHeaders.AddBlockLocatorHash(&btcwire.GenesisHash)
	}
	maxGetHeaders.BlockLocatorHashes = append(maxGetHeaders.BlockLocatorHashes,
		&btcwire.GenesisHash)
	//maxGetHeaders.InvList = append(maxGetData.InvList, iv)
	maxGetHeadersEncoded := []byte{
		0x62, 0xea, 0x00, 0x00, // Protocol version 60002
		0xfd, 0xf5, 0x01, // Varint for number of block loc hashes (501)
	}

	tests := []struct {
		in       *btcwire.MsgGetHeaders // Value to encode
		buf      []byte                 // Wire encoding
		pver     uint32                 // Protocol version for wire encoding
		max      int                    // Max size of fixed buffer to induce errors
		writeErr error                  // Expected write error
		readErr  error                  // Expected read error
	}{
		// Force error in protocol version.
		{baseGetHeaders, baseGetHeadersEncoded, pver, 0, io.ErrShortWrite, io.EOF},
		// Force error in block locator hash count.
		{baseGetHeaders, baseGetHeadersEncoded, pver, 4, io.ErrShortWrite, io.EOF},
		// Force error in block locator hashes.
		{baseGetHeaders, baseGetHeadersEncoded, pver, 5, io.ErrShortWrite, io.EOF},
		// Force error in stop hash.
		{baseGetHeaders, baseGetHeadersEncoded, pver, 69, io.ErrShortWrite, io.EOF},
		// Force error with greater than max block locator hashes.
		{maxGetHeaders, maxGetHeadersEncoded, pver, 7, btcwireErr, btcwireErr},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode to wire format.
		w := newFixedWriter(test.max)
		err := test.in.BtcEncode(w, test.pver)
		if reflect.TypeOf(err) != reflect.TypeOf(test.writeErr) {
			t.Errorf("BtcEncode #%d wrong error got: %v, want: %v",
				i, err, test.writeErr)
			continue
		}

		// For errors which are not of type btcwire.MessageError, check
		// them for equality.
		if _, ok := err.(*btcwire.MessageError); !ok {
			if err != test.writeErr {
				t.Errorf("BtcEncode #%d wrong error got: %v, "+
					"want: %v", i, err, test.writeErr)
				continue
			}
		}

		// Decode from wire format.
		var msg btcwire.MsgGetHeaders
		r := newFixedReader(test.max, test.buf)
		err = msg.BtcDecode(r, test.pver)
		if reflect.TypeOf(err) != reflect.TypeOf(test.readErr) {
			t.Errorf("BtcDecode #%d wrong error got: %v, want: %v",
				i, err, test.readErr)
			continue
		}

		// For errors which are not of type btcwire.MessageError, check
		// them for equality.
		if _, ok := err.(*btcwire.MessageError); !ok {
			if err != test.readErr {
				t.Errorf("BtcEncode #%d wrong error got: %v, "+
					"want: %v", i, err, test.readErr)
				continue
			}
		}
	}
}
<|endoftext|>


<commit_before>package audio

import (
	"bitbucket.org/oakmoundstudio/oak/dlog"
	"bitbucket.org/oakmoundstudio/oak/physics"
)

func (a *Audio) Play() error {
	var err error
	ad := a.AudioData
	if a.F.Volume != 0 {
		err = a.ShiftVolume(a.F.Volume)
		if err != nil {
			return err
		}
	}
	if a.F.Frequency != 0 {
		err = a.ShiftFrequency(a.F.Frequency)
		if err != nil {
			return err
		}
	}
	switch a.F.ForceLoop {
	case FORCE_LOOP:
		ad.SetLooping(true)
	case FORCE_NO_LOOP:
		ad.SetLooping(false)
	}

	if a.F.Ears != nil && a.X != nil && a.Y != nil {
		err = a.ScaleVolume(a.F.CalculateVolume(physics.NewVector(*a.X, *a.Y)))
		if err != nil {
			return err
		}
		err = a.SetPan(a.F.CalculatePan(*a.X))
		if err != nil {
			return err
		}
	} else {
		err = a.ShiftPan(a.F.Pan)
		if err != nil {
			return err
		}
	}

	return ad.Play()
}

func (f *Font) PlayWav(filename string) error {
	ad, err := GetWav(filename)
	if err == nil {
		a := Audio{ad, f, nil, nil}
		err = a.Play()
	} else {
		dlog.Error(err)
	}
	return err
}

func PlayWav(filename string) error {
	return defFont.PlayWav(filename)
}
<commit_msg>Fix audio constantly getting quieter<commit_after>package audio

import (
	"bitbucket.org/oakmoundstudio/oak/dlog"
	"bitbucket.org/oakmoundstudio/oak/physics"
)

func (a *Audio) Play() error {
	origVolume, _ := a.GetVolume()
	origFrequency, _ := a.GetFrequency()
	origPan, _ := a.GetPan()
	var err error
	ad := a.AudioData
	if a.F.Volume != 0 {
		err = a.ShiftVolume(a.F.Volume)
		if err != nil {
			return err
		}
	}
	if a.F.Frequency != 0 {
		err = a.ShiftFrequency(a.F.Frequency)
		if err != nil {
			return err
		}
	}
	switch a.F.ForceLoop {
	case FORCE_LOOP:
		ad.SetLooping(true)
	case FORCE_NO_LOOP:
		ad.SetLooping(false)
	}

	if a.F.Ears != nil && a.X != nil && a.Y != nil {
		err = a.ScaleVolume(a.F.CalculateVolume(physics.NewVector(*a.X, *a.Y)))
		if err != nil {
			return err
		}
		err = a.SetPan(a.F.CalculatePan(*a.X))
		if err != nil {
			return err
		}
	} else {
		err = a.ShiftPan(a.F.Pan)
		if err != nil {
			return err
		}
	}

	err = ad.Play()
	ad.SetVolume(origVolume)
	ad.SetFrequency(origFrequency)
	ad.SetPan(origPan)
	return err
}

func (f *Font) PlayWav(filename string) error {
	ad, err := GetWav(filename)
	if err == nil {
		a := Audio{ad, f, nil, nil}
		err = a.Play()
	} else {
		dlog.Error(err)
	}
	return err
}

func PlayWav(filename string) error {
	return defFont.PlayWav(filename)
}
<|endoftext|>


<commit_before>package auth2

import (
	"database/sql"
	"time"

	"github.com/lib/pq"
)

type user struct {
	id                string
	createdAt         time.Time
	appId             string
	email             string
	hashedPass        string
	lang              string
	confirmationKey   string
	confirmationKeyAt time.Time
	resetKey          string
	resetKeyAt        time.Time
}

type store interface {
	createSchema() error

	createUser(userId string, createdAt time.Time, appId, email, hashedPass, lang, confirmationKey string) error
	setUserConfirmationKeyAt(appId, email string, confirmationKeyAt time.Time) error
	setUserResetKey(appId, email, resetKey string, resetKeyAt time.Time) error
	setUserHashedPass(appId, email, hashedPass string) error

	getUserId(appId, email string) (userId string, err error)
	getUser(userId string) (user user, err error)

	getUserConfirmation(appId, email string) (confirmationKey string, confirmationKeyAt time.Time, err error)
	getUserPassword(appId, email string) (userId, hashedPass string, confirmationKeyAt time.Time, err error)
	// getUserResetKey(appId, email string) (resetKey string, resetKeyAt time.Time, err error)

	createSession(id, userId string, createdAt time.Time) error
	removeSession(id string) error
	getSession(sessionId string) (userId string, createdAt time.Time, err error)
}

type storePg struct {
	db *sql.DB
}

func NewStorePg(db *sql.DB) storePg {
	return storePg{db}
}

func (self storePg) createSchema() error {
	schema := `
		CREATE SCHEMA auth;

		CREATE TYPE auth.lang AS ENUM ('pt_PT', 'en_US');

		CREATE TABLE auth.user (
		   id                CHAR(36) NOT NULL,
		   createdAt         TIMESTAMPTZ NOT NULL,
		   appId             TEXT NOT NULL,
		   email             TEXT NOT NULL,
		   hashedPass        TEXT NOT NULL,
		   lang              auth.lang NOT NULL,
		   confirmationKey   CHAR(36) NOT NULL,
		   confirmationKeyAt TIMESTAMPTZ,
		   resetKey          CHAR(36),
		   resetKeyAt        TIMESTAMPTZ,

		   CONSTRAINT pk_auth_user PRIMARY KEY (id)
		);

		CREATE UNIQUE INDEX idx_auth_user_appId_email ON auth.user (appId, email);

		CREATE TABLE auth.session (
			id         CHAR(36) NOT NULL,
			userId     CHAR(36) NOT NULL,
			createdAt  TIMESTAMPTZ NOT NULL,

			CONSTRAINT pk_auth_session PRIMARY KEY (id),
			CONSTRAINT fk_auth_session_userId FOREIGN KEY (userId) REFERENCES auth.user(id)
		);
	`

	_, err := self.db.Exec(schema)
	return err
}

func (self storePg) createUser(userId string, createdAt time.Time, appId, email, hashedPass, lang, confirmationKey string) error {
	insert := `
		INSERT INTO auth.user
		(id, createdAt, appId, email, hashedPass, lang, confirmationKey)
		VALUES
		($1, $2, $3, $4, $5, $6, $7);
	`

	stmt, err := self.db.Prepare(insert)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(userId, createdAt, appId, email, hashedPass, lang, confirmationKey)
	return err
}

func (self storePg) setUserConfirmationKeyAt(appId, email string, confirmationKeyAt time.Time) error {
	update := `
		UPDATE auth.user
		SET confirmationKeyAt = $1
		WHERE appId = $2 AND email = $3;
	`

	stmt, err := self.db.Prepare(update)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(confirmationKeyAt, appId, email)
	return err
}

func (self storePg) setUserResetKey(appId, email, resetKey string, resetKeyAt time.Time) error {
	update := `
		UPDATE auth.user
		SET resetKey = $1, resetKeyAt = $2
		WHERE appId = $3 AND email = $4;
	`

	stmt, err := self.db.Prepare(update)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(resetKey, resetKeyAt, appId, email)
	return err
}

func (self storePg) setUserHashedPass(appId, email, hashedPass string) error {
	update := `
		UPDATE auth.user
		SET hashedPass = $1, resetKey = NULL, resetKeyAt = NULL
		WHERE appId = $2 AND email = $3;
	`

	stmt, err := self.db.Prepare(update)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(hashedPass, appId, email)
	return err
}

func (self storePg) getUserId(appId, email string) (userId string, err error) {
	query := `
		SELECT id
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&userId)
	return
}

func (self storePg) getUser(userId string) (user user, err error) {
	var scanConfirmationKeyAt pq.NullTime
	var scanResetKey sql.NullString
	var scanResetKeyAt pq.NullTime
	query := `
		SELECT createdAt, appId, email, hashedPass, lang, confirmationKey, confirmationKeyAt, resetKey, resetKeyAt
		FROM auth.user
		WHERE id = $1;
	`
	err = self.db.QueryRow(query, userId).Scan(
		&user.createdAt,
		&user.appId,
		&user.email,
		&user.hashedPass,
		&user.lang,
		&user.confirmationKey,
		&scanConfirmationKeyAt,
		&scanResetKey,
		&scanResetKeyAt,
	)

	if scanConfirmationKeyAt.Valid {
		user.confirmationKeyAt = scanConfirmationKeyAt.Time
	}
	if scanResetKey.Valid {
		user.resetKey = scanResetKey.String
	}
	if scanResetKeyAt.Valid {
		user.resetKeyAt = scanResetKeyAt.Time
	}
	return
}

func (self storePg) getUserConfirmation(appId, email string) (confirmationKey string, confirmationKeyAt time.Time, err error) {
	var scanConfirmationKeyAt pq.NullTime
	query := `
		SELECT confirmationKey, confirmationKeyAt
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&confirmationKey, &scanConfirmationKeyAt)
	if scanConfirmationKeyAt.Valid {
		confirmationKeyAt = scanConfirmationKeyAt.Time
	}
	return
}

func (self storePg) getUserPassword(appId, email string) (userId, hashedPass string, confirmationKeyAt time.Time, err error) {
	query := `
		SELECT id, hashedPass, confirmationKeyAt
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&userId, &hashedPass, &confirmationKeyAt)
	return
}

func (self storePg) getUserResetKey(appId, email string) (resetKey string, resetKeyAt time.Time, err error) {
	var scanResetKey sql.NullString
	var scanSetResetKeyAt pq.NullTime
	query := `
		SELECT resetKey, resetKeyAt
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&scanResetKey, &scanSetResetKeyAt)
	if scanResetKey.Valid {
		resetKey = scanResetKey.String
	}
	if scanSetResetKeyAt.Valid {
		resetKeyAt = scanSetResetKeyAt.Time
	}
	return
}

func (self storePg) createSession(id, userId string, createdAt time.Time) error {
	insert := `
		INSERT INTO auth.session
		(id, userId, createdAt)
		VALUES
		($1, $2, $3);
	`

	stmt, err := self.db.Prepare(insert)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(id, userId, createdAt)
	return err
}

func (self storePg) removeSession(id string) error {
	delete := `
		DELETE FROM auth.session
		WHERE id = $1;
	`

	stmt, err := self.db.Prepare(delete)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(id)
	return err
}

func (self storePg) getSession(sessionId string) (userId string, createdAt time.Time, err error) {
	query := `
		SELECT userId, createdAt
		FROM auth.session
		WHERE id = $1;
	`
	err = self.db.QueryRow(query, sessionId).Scan(&userId, &createdAt)
	return
}
<commit_msg>update<commit_after>package auth2

import (
	"database/sql"
	"time"

	"github.com/lib/pq"
)

type user struct {
	id                string
	createdAt         time.Time
	appId             string
	email             string
	hashedPass        string
	lang              string
	confirmationKey   string
	confirmationKeyAt time.Time
	resetKey          string
	resetKeyAt        time.Time
}

type store interface {
	createSchema() error

	createUser(userId string, createdAt time.Time, appId, email, hashedPass, lang, confirmationKey string) error
	setUserConfirmationKeyAt(appId, email string, confirmationKeyAt time.Time) error
	setUserResetKey(appId, email, resetKey string, resetKeyAt time.Time) error
	setUserHashedPass(appId, email, hashedPass string) error

	getUserId(appId, email string) (userId string, err error)
	getUser(userId string) (user user, err error)

	getUserConfirmation(appId, email string) (confirmationKey string, confirmationKeyAt time.Time, err error)
	getUserPassword(appId, email string) (userId, hashedPass string, confirmationKeyAt time.Time, err error)
	// getUserResetKey(appId, email string) (resetKey string, resetKeyAt time.Time, err error)

	createSession(id, userId string, createdAt time.Time) error
	removeSession(id string) error
	getSession(sessionId string) (userId string, createdAt time.Time, err error)
}

type storePg struct {
	db *sql.DB
}

func NewStorePg(db *sql.DB) storePg {
	return storePg{db}
}

func (self storePg) createSchema() error {
	schema := `
		CREATE SCHEMA auth;

		CREATE TYPE auth.lang AS ENUM ('pt_PT', 'en_US');

		CREATE TABLE auth.user (
		   id                CHAR(36) NOT NULL,
		   createdAt         TIMESTAMPTZ NOT NULL,
		   appId             TEXT NOT NULL,
		   email             TEXT NOT NULL,
		   hashedPass        TEXT NOT NULL,
		   lang              auth.lang NOT NULL,
		   confirmationKey   CHAR(36) NOT NULL,
		   confirmationKeyAt TIMESTAMPTZ,
		   resetKey          CHAR(36),
		   resetKeyAt        TIMESTAMPTZ,

		   CONSTRAINT pk_auth_user PRIMARY KEY (id)
		);

		CREATE UNIQUE INDEX idx_auth_user_appId_email ON auth.user (appId, email);

		CREATE TABLE auth.session (
			id         CHAR(36) NOT NULL,
			userId     CHAR(36) NOT NULL,
			createdAt  TIMESTAMPTZ NOT NULL,

			CONSTRAINT pk_auth_session PRIMARY KEY (id),
			CONSTRAINT fk_auth_session_userId FOREIGN KEY (userId) REFERENCES auth.user(id)
		);
	`

	_, err := self.db.Exec(schema)
	return err
}

func (self storePg) createUser(userId string, createdAt time.Time, appId, email, hashedPass, lang, confirmationKey string) error {
	insert := `
		INSERT INTO auth.user
		(id, createdAt, appId, email, hashedPass, lang, confirmationKey)
		VALUES
		($1, $2, $3, $4, $5, $6, $7);
	`

	stmt, err := self.db.Prepare(insert)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(userId, createdAt, appId, email, hashedPass, lang, confirmationKey)
	return err
}

func (self storePg) setUserConfirmationKeyAt(appId, email string, confirmationKeyAt time.Time) error {
	update := `
		UPDATE auth.user
		SET confirmationKeyAt = $1
		WHERE appId = $2 AND email = $3;
	`

	stmt, err := self.db.Prepare(update)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(confirmationKeyAt, appId, email)
	return err
}

func (self storePg) setUserResetKey(appId, email, resetKey string, resetKeyAt time.Time) error {
	update := `
		UPDATE auth.user
		SET resetKey = $1, resetKeyAt = $2
		WHERE appId = $3 AND email = $4;
	`

	stmt, err := self.db.Prepare(update)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(resetKey, resetKeyAt, appId, email)
	return err
}

func (self storePg) setUserHashedPass(appId, email, hashedPass string) error {
	update := `
		UPDATE auth.user
		SET hashedPass = $1, resetKey = NULL, resetKeyAt = NULL
		WHERE appId = $2 AND email = $3;
	`

	stmt, err := self.db.Prepare(update)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(hashedPass, appId, email)
	return err
}

func (self storePg) getUserId(appId, email string) (userId string, err error) {
	query := `
		SELECT id
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&userId)
	return
}

func (self storePg) getUser(userId string) (user user, err error) {
	query := `
		SELECT createdAt, appId, email, hashedPass, lang, confirmationKey, confirmationKeyAt, resetKey, resetKeyAt
		FROM auth.user
		WHERE id = $1;
	`

	var scanConfirmationKeyAt pq.NullTime
	var scanResetKey sql.NullString
	var scanResetKeyAt pq.NullTime

	err = self.db.QueryRow(query, userId).Scan(
		&user.createdAt,
		&user.appId,
		&user.email,
		&user.hashedPass,
		&user.lang,
		&user.confirmationKey,
		&scanConfirmationKeyAt,
		&scanResetKey,
		&scanResetKeyAt,
	)

	if scanConfirmationKeyAt.Valid {
		user.confirmationKeyAt = scanConfirmationKeyAt.Time
	}
	if scanResetKey.Valid {
		user.resetKey = scanResetKey.String
	}
	if scanResetKeyAt.Valid {
		user.resetKeyAt = scanResetKeyAt.Time
	}

	return
}

func (self storePg) getUserConfirmation(appId, email string) (confirmationKey string, confirmationKeyAt time.Time, err error) {
	var scanConfirmationKeyAt pq.NullTime
	query := `
		SELECT confirmationKey, confirmationKeyAt
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&confirmationKey, &scanConfirmationKeyAt)
	if scanConfirmationKeyAt.Valid {
		confirmationKeyAt = scanConfirmationKeyAt.Time
	}
	return
}

func (self storePg) getUserPassword(appId, email string) (userId, hashedPass string, confirmationKeyAt time.Time, err error) {
	query := `
		SELECT id, hashedPass, confirmationKeyAt
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&userId, &hashedPass, &confirmationKeyAt)
	return
}

func (self storePg) getUserResetKey(appId, email string) (resetKey string, resetKeyAt time.Time, err error) {
	var scanResetKey sql.NullString
	var scanSetResetKeyAt pq.NullTime
	query := `
		SELECT resetKey, resetKeyAt
		FROM auth.user
		WHERE appId = $1 AND email = $2;
	`
	err = self.db.QueryRow(query, appId, email).Scan(&scanResetKey, &scanSetResetKeyAt)
	if scanResetKey.Valid {
		resetKey = scanResetKey.String
	}
	if scanSetResetKeyAt.Valid {
		resetKeyAt = scanSetResetKeyAt.Time
	}
	return
}

func (self storePg) createSession(id, userId string, createdAt time.Time) error {
	insert := `
		INSERT INTO auth.session
		(id, userId, createdAt)
		VALUES
		($1, $2, $3);
	`

	stmt, err := self.db.Prepare(insert)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(id, userId, createdAt)
	return err
}

func (self storePg) removeSession(id string) error {
	delete := `
		DELETE FROM auth.session
		WHERE id = $1;
	`

	stmt, err := self.db.Prepare(delete)
	if err != nil {
		return err
	}

	_, err = stmt.Exec(id)
	return err
}

func (self storePg) getSession(sessionId string) (userId string, createdAt time.Time, err error) {
	query := `
		SELECT userId, createdAt
		FROM auth.session
		WHERE id = $1;
	`
	err = self.db.QueryRow(query, sessionId).Scan(&userId, &createdAt)
	return
}
<|endoftext|>


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

package auth

import (
	"io/ioutil"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fsnotify/fsnotify"
)

const (
	// expirationDelta is a margin of error during which a token should be
	// considered expired. This should help avoid expiration races when server
	// times don't match
	expirationDelta = 10 * time.Second
)

var (
	// TokenFileName is the file in which we store the current token
	TokenFileName = "auth.token"

	currentToken    string
	currentIdentity Identity
	zerotime        time.Time
	expiration      time.Time
	cond            = &sync.Cond{L: &sync.Mutex{}}
)

// TokenFunc is a function that can return an authentication token and its
// expiration time
type TokenFunc func() (string, int64, error)

// RefreshToken gets a new token, sets it as the current, and returns the expiration time
func RefreshToken(f TokenFunc, filename string) (int64, error) {
	log.Debug("Refreshing authentication token")
	token, expires, err := f()
	if err != nil {
		return 0, err
	}
	updateToken(token, time.Unix(expires, 0), filename)
	log.WithField("expiration", expires).Info("Received new authentication token")
	return expires, err
}

// AuthToken returns an unexpired auth token, blocking if necessary until
// authenticated
func AuthToken() string {
	cond.L.Lock()
	defer cond.L.Unlock()
	for expired() {
		cond.Wait()
	}
	return currentToken
}

// A non-blocking call to get an unexpired auth token.  Returns an error
//  If no token exists or if the token is expired
func AuthTokenNonBlocking() (string, error) {
	if currentToken == "" {
		return "", ErrNotAuthenticated
	}

	if expired() {
		return "", ErrIdentityTokenExpired
	}

	return currentToken, nil
}

// MasterToken() generates a new token with an empty host and pool ID and the master's public key,
//  signed by the master's private key.  This will return an error if there is no master private
//  key available (i.e. if we are not the master)
func MasterToken() (string, error) {
	masterpublic, err := GetMasterPublicKey()
	if err != nil {
		return "", err
	}
	keypem, err := PEMFromRSAPublicKey(masterpublic, nil)
	if err != nil {
		return "", err
	}

	signed, _, err := CreateJWTIdentity("", "", true, true, keypem, time.Hour)
	if err != nil {
		return "", err
	}

	return signed, nil
}

// CurrentIdentity returns the identity represented by the currently-live token,
// or nil if the token is not yet available
func CurrentIdentity() Identity {
	cond.L.Lock()
	defer cond.L.Unlock()
	return currentIdentity
}

// TokenLoop accepts a function that returns an expiring token. It will then
// periodically refresh that token, one minute before it is due to expire,
// setting the result as the current live token, until the done channel is
// closed.
func TokenLoop(f TokenFunc, tokenfile string, done <-chan interface{}) {
	for {
		expires, err := RefreshToken(f, tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to obtain authentication token. Retrying in 10s")
			select {
			case <-done:
				return
			case <-time.After(10 * time.Second):
			}
			continue
		}
		// Reauthenticate 1 minute before the token expires
		expiration := time.Unix(expires, 0).Sub(now())
		refresh := expiration - time.Duration(1*time.Minute)
		select {
		case <-done:
			return
		case <-time.After(refresh):
		}
	}
}

// Watch a token file for changes. Load the token when those changes occur.
func WatchTokenFile(tokenfile string, done <-chan interface{}) error {
	log := log.WithFields(logrus.Fields{
		"tokenfile": tokenfile,
	})

	loadToken := func() {
		data, err := ioutil.ReadFile(tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to load authentication token from file. Continuing to watch for changes")
		}
		// No need to handle expires or save file, because we're loading from the file rather
		// than re-requesting authentication tokens
		updateToken(string(data), zerotime, "")
		log.Infof("Updated authentication token from disk")
	}

	// An initial token load without any file changes
	loadToken()

	// Now watch for changes
	filechangechan, err := NotifyOnChange(tokenfile, fsnotify.Write|fsnotify.Create, done)
	if err != nil {
		return err
	}
	for _ = range filechangechan {
		loadToken()
	}
	return nil
}

func now() time.Time {
	return time.Now().UTC()
}

// expired returns whether the currently-live token has expired. If the token
// is empty, is is considered expired for these purposes. If it is the zero
// instant, it never expires. Otherwise, it is expired if the expiration time
// minus a margin of error is in the past.
func expired() bool {
	if currentToken == "" {
		return true
	}
	if expiration.IsZero() {
		return false
	}
	return expiration.Add(-expirationDelta).Before(now())
}

func updateToken(token string, expires time.Time, filename string) {
	WaitForDelegateKeys()
	cond.L.Lock()
	currentToken = token
	currentIdentity = getIdentityFromToken(token)
	expiration = expires
	if filename != "" {
		ioutil.WriteFile(filename, []byte(token), 0600)
	}
	cond.L.Unlock()
	cond.Broadcast()
}

func getIdentityFromToken(token string) Identity {
	identity, err := ParseJWTIdentity(token)
	if err != nil {
		log.WithError(err).Error("Unable to obtain identity from token.")
		return nil
	}
	return identity
}
<commit_msg>Flipped one log message to debug<commit_after>// Copyright 2016 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package auth

import (
	"io/ioutil"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fsnotify/fsnotify"
)

const (
	// expirationDelta is a margin of error during which a token should be
	// considered expired. This should help avoid expiration races when server
	// times don't match
	expirationDelta = 10 * time.Second
)

var (
	// TokenFileName is the file in which we store the current token
	TokenFileName = "auth.token"

	currentToken    string
	currentIdentity Identity
	zerotime        time.Time
	expiration      time.Time
	cond            = &sync.Cond{L: &sync.Mutex{}}
)

// TokenFunc is a function that can return an authentication token and its
// expiration time
type TokenFunc func() (string, int64, error)

// RefreshToken gets a new token, sets it as the current, and returns the expiration time
func RefreshToken(f TokenFunc, filename string) (int64, error) {
	log.Debug("Refreshing authentication token")
	token, expires, err := f()
	if err != nil {
		return 0, err
	}
	updateToken(token, time.Unix(expires, 0), filename)
	log.WithField("expiration", expires).Debug("Received new authentication token")
	return expires, err
}

// AuthToken returns an unexpired auth token, blocking if necessary until
// authenticated
func AuthToken() string {
	cond.L.Lock()
	defer cond.L.Unlock()
	for expired() {
		cond.Wait()
	}
	return currentToken
}

// A non-blocking call to get an unexpired auth token.  Returns an error
//  If no token exists or if the token is expired
func AuthTokenNonBlocking() (string, error) {
	if currentToken == "" {
		return "", ErrNotAuthenticated
	}

	if expired() {
		return "", ErrIdentityTokenExpired
	}

	return currentToken, nil
}

// MasterToken() generates a new token with an empty host and pool ID and the master's public key,
//  signed by the master's private key.  This will return an error if there is no master private
//  key available (i.e. if we are not the master)
func MasterToken() (string, error) {
	masterpublic, err := GetMasterPublicKey()
	if err != nil {
		return "", err
	}
	keypem, err := PEMFromRSAPublicKey(masterpublic, nil)
	if err != nil {
		return "", err
	}

	signed, _, err := CreateJWTIdentity("", "", true, true, keypem, time.Hour)
	if err != nil {
		return "", err
	}

	return signed, nil
}

// CurrentIdentity returns the identity represented by the currently-live token,
// or nil if the token is not yet available
func CurrentIdentity() Identity {
	cond.L.Lock()
	defer cond.L.Unlock()
	return currentIdentity
}

// TokenLoop accepts a function that returns an expiring token. It will then
// periodically refresh that token, one minute before it is due to expire,
// setting the result as the current live token, until the done channel is
// closed.
func TokenLoop(f TokenFunc, tokenfile string, done <-chan interface{}) {
	for {
		expires, err := RefreshToken(f, tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to obtain authentication token. Retrying in 10s")
			select {
			case <-done:
				return
			case <-time.After(10 * time.Second):
			}
			continue
		}
		// Reauthenticate 1 minute before the token expires
		expiration := time.Unix(expires, 0).Sub(now())
		refresh := expiration - time.Duration(1*time.Minute)
		select {
		case <-done:
			return
		case <-time.After(refresh):
		}
	}
}

// Watch a token file for changes. Load the token when those changes occur.
func WatchTokenFile(tokenfile string, done <-chan interface{}) error {
	log := log.WithFields(logrus.Fields{
		"tokenfile": tokenfile,
	})

	loadToken := func() {
		data, err := ioutil.ReadFile(tokenfile)
		if err != nil {
			log.WithError(err).Warn("Unable to load authentication token from file. Continuing to watch for changes")
		}
		// No need to handle expires or save file, because we're loading from the file rather
		// than re-requesting authentication tokens
		updateToken(string(data), zerotime, "")
		log.Infof("Updated authentication token from disk")
	}

	// An initial token load without any file changes
	loadToken()

	// Now watch for changes
	filechangechan, err := NotifyOnChange(tokenfile, fsnotify.Write|fsnotify.Create, done)
	if err != nil {
		return err
	}
	for _ = range filechangechan {
		loadToken()
	}
	return nil
}

func now() time.Time {
	return time.Now().UTC()
}

// expired returns whether the currently-live token has expired. If the token
// is empty, is is considered expired for these purposes. If it is the zero
// instant, it never expires. Otherwise, it is expired if the expiration time
// minus a margin of error is in the past.
func expired() bool {
	if currentToken == "" {
		return true
	}
	if expiration.IsZero() {
		return false
	}
	return expiration.Add(-expirationDelta).Before(now())
}

func updateToken(token string, expires time.Time, filename string) {
	WaitForDelegateKeys()
	cond.L.Lock()
	currentToken = token
	currentIdentity = getIdentityFromToken(token)
	expiration = expires
	if filename != "" {
		ioutil.WriteFile(filename, []byte(token), 0600)
	}
	cond.L.Unlock()
	cond.Broadcast()
}

func getIdentityFromToken(token string) Identity {
	identity, err := ParseJWTIdentity(token)
	if err != nil {
		log.WithError(err).Error("Unable to obtain identity from token.")
		return nil
	}
	return identity
}
<|endoftext|>


<commit_before>package zk

/**
This test depends on a running zookeeper instance locally. If you cannot
run a zookeeper instance locally, please disable this test.

How to run zookeeper locally:

(1) Install a zookeeper package;

(2) Create a data dir for zookeeper:

mk -p /tmp/cfg/zk

(3) Create a config file with following text:

tickTime=2000
dataDir=/tmp/cfg/zk
clientPort=18888

(4) Run zookeeper:

./zkServer.sh start <Your zookeeper config file path>
*/

import (
	"testing"
	"fmt"
)

const (
	// Zookeeper port, pls change this to be the real port number
	// on your host.
	Port = 18888
	// Zookeeper timeout value.
	Timeout = 4000
)

func TestZkCreate(t *testing.T) {
	path := "/testZkCreate"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	statRes := h.Exists(path)
	if statRes.GetRc() != ZOK {
		t.Error("Fails to call exists()")
	}
}

func TestZkGet(t *testing.T) {
	path := "/testZkGet"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	dataRes := h.Get(path)
	if dataRes.GetRc() != ZOK {
		t.Error("Fails to call exists()")
	}

	if value != string(dataRes.GetData()) {
		t.Error("Fails to get the same data!")
	}
}

func TestZkExistW(t *testing.T) {
	path := "/testZkExistW"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	_,watcher := h.ExistsW(path)

	// Test channel is not available.
	select {
		case _,ok := <-watcher:
			if ok {
				t.Error("Value should not be available yet")
			} else {
				t.Error("Channel should not be closed yet")
			}
		// Expect default behavior
		default:
	}

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	// Test channel is available by now.
	select {
		case event,ok := <-watcher:
			if ok {
				if event.Type != ZOO_CREATED_EVENT {
					t.Error("Expect created event")
				}
				if event.State != ZOO_CONNECTED_STATE {
					t.Error("Expect connected state")
				}
				if event.Path != path {
					t.Error("Path mismatch")
				}
			} else {
				t.Error("Channel should not be closed yet")
			}
		// Expect default behavior
		default:
			t.Error("Channel should be available by now")
	}
}
<commit_msg>Adding a few test cases<commit_after>package zk

/**
This test depends on a running zookeeper instance locally. If you cannot
run a zookeeper instance locally, please disable this test.

How to run zookeeper locally:

(1) Install a zookeeper package;

(2) Create a data dir for zookeeper:

mk -p /tmp/cfg/zk

(3) Create a config file with following text:

tickTime=2000
dataDir=/tmp/cfg/zk
clientPort=18888

(4) Run zookeeper:

./zkServer.sh start <Your zookeeper config file path>
*/

import (
	"testing"
	"fmt"
)

const (
	// Zookeeper port, pls change this to be the real port number
	// on your host.
	Port = 18888
	// Zookeeper timeout value.
	Timeout = 4000
)

func TestZkCreate(t *testing.T) {
	path := "/testZkCreate"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	statRes := h.Exists(path)
	if statRes.GetRc() != ZOK {
		t.Error("Fails to call exists()")
	}

	if statRes.GetVersion() != 0 {
		t.Error("Expected version ", statRes.GetVersion())
	}
}

func TestZkGet(t *testing.T) {
	path := "/testZkGet"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	dataRes := h.Get(path)
	if dataRes.GetRc() != ZOK {
		t.Error("Fails to call exists()")
	}

	if value != string(dataRes.GetData()) {
		t.Error("Fails to get the same data!")
	}
}

func TestZkExistW(t *testing.T) {
	path := "/testZkExistW"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	_,watcher := h.ExistsW(path)

	// Test channel is not available.
	select {
		case _,ok := <-watcher:
			if ok {
				t.Error("Value should not be available yet")
			} else {
				t.Error("Channel should not be closed yet")
			}
		// Expect default behavior
		default:
	}

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	// Test channel is available by now.
	select {
		case event,ok := <-watcher:
			if ok {
				if event.Type != ZOO_CREATED_EVENT {
					t.Error("Expect created event")
				}
				if event.State != ZOO_CONNECTED_STATE {
					t.Error("Expect connected state")
				}
				if event.Path != path {
					t.Error("Path mismatch")
				}
			} else {
				t.Error("Channel should not be closed yet")
			}
		// Expect default behavior
		default:
			t.Error("Channel should be available by now")
	}

	// Verify that data has been written correctly.
	dataRes := h.Get(path)
	if dataRes.GetRc() != ZOK {
		t.Error("Fails to call exists()")
	}

	if value != string(dataRes.GetData()) {
		t.Error("Fails to get the same data!")
	}
}

func TestZkSet(t *testing.T) {
	path := "/testZkSet"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path, -1)

	// Create a new znode.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	newValue := "newValue"
	statRes := h.Set(path, []byte(newValue), 0)

	if statRes.GetRc() != ZOK {
		t.Error("Fails to set a value")
	}
	if statRes.GetVersion() != 1 {
		t.Error("Real version number is ", statRes.GetVersion())
	}
}

func TestZkGetChildren(t *testing.T) {
	path := "/testZkGetChildren"
	value := "something"
	service := fmt.Sprintf("localhost:%d", Port)
	h,ok := NewZHandle(service, Timeout, nil)
	if !ok {
		t.Error("Fails to create a zookeeper handle")
	}

	// Make sure that previous leftovers are all gone.
	h.Delete(path + "/a", -1)
	h.Delete(path + "/b", -1)
	h.Delete(path, -1)

	// Create new znodes.
	strRes := h.Create(path, value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	strRes = h.Create(path + "/a", value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	strRes = h.Create(path + "/b", value, ZOO_OPEN_ACLS, 0)
	if strRes.GetRc() != ZOK {
		t.Error("Fails to create a path")
	}

	// Verify GetChildren() is correct.
	stringsRes := h.GetChildren(path)
	if stringsRes.GetRc() != ZOK {
		t.Error("Fails to call GetChildren")
	}
	strings := stringsRes.GetStrings()
	if len(strings) != 2 {
		t.Error("Has number of strings: ", len(strings))
	}
}
<|endoftext|>


<commit_before>package workloads

import (
	"bytes"
	"log"
	"math"
	"math/rand"
	"strconv"
	"sync"
	"time"

	"github.com/pavel-paulau/blurr/databases"
)

type DefaultWorkload struct {
	Config       Config
	DeletedItems int64
}

func (workload *DefaultWorkload) Init(config Config) {
	workload.Config = config
}

// Generate new *unique* key
func (workload *DefaultWorkload) GenerateNewKey(currentRecords int64) string {
	strCurrentRecords := strconv.FormatInt(currentRecords, 10)
	return Hash(strCurrentRecords)
}

// Generate random key from current key space
func (workload *DefaultWorkload) GenerateExistingKey(currentRecords int64) string {
	rand.Seed(time.Now().UnixNano())
	randRecord := workload.DeletedItems + rand.Int63n(currentRecords-workload.DeletedItems)
	strRandRecord := strconv.FormatInt(randRecord, 10)
	return Hash(strRandRecord)
}

// Generate sequential key for removal
func (workload *DefaultWorkload) GenerateKeyForRemoval() string {
	keyForRemoval := strconv.FormatInt(workload.DeletedItems+1, 10)
	workload.DeletedItems++
	return Hash(keyForRemoval)
}

// Generate value with deterministic indexable fields and arbitrary body
func (workload *DefaultWorkload) GenerateValue(key string,
	indexableFields, size int) map[string]interface{} {

	// Hex lengh is 32 characters, so only 22 indexable fields are allowed
	if indexableFields >= 20 {
		log.Fatal("Too much fields! It must be less than 20")
	}
	// Gererate indexable fields (shifting over key name)
	value := map[string]interface{}{}
	for i := 0; i < indexableFields; i++ {
		fieldName := "field" + strconv.Itoa(i)
		value[fieldName] = fieldName + "-" + key[i:i+10]
	}
	// Generate value body in order to meet value size specification
	fieldName := "field" + strconv.Itoa(indexableFields)
	buffer := bytes.Buffer{}
	bodyHash := Hash(key)
	iterations := (size - len(fieldName+"-"+key[:10])*indexableFields) / 32
	for i := 0; i < iterations; i++ {
		buffer.WriteString(bodyHash)
	}
	value[fieldName] = buffer.String()
	return value
}

func (workload *DefaultWorkload) GenerateQuery(indexableFields int,
	currentRecords int64) (string, string, int) {

	i := rand.Intn(indexableFields)
	fieldName := "field" + strconv.Itoa(i)
	fieldValue := fieldName + "-" + workload.GenerateExistingKey(currentRecords)[i:i+10]
	limit := 10 + rand.Intn(10)
	return fieldName, fieldValue, limit
}

// Generate slice of shuffled characters (CRUD-Q shorthands)
func (workload *DefaultWorkload) PrepareBatch() []string {
	operations := make([]string, 0, 100)
	randOperations := make([]string, 100, 100)
	for i := 0; i < workload.Config.CreatePercentage; i++ {
		operations = append(operations, "c")
	}
	for i := 0; i < workload.Config.ReadPercentage; i++ {
		operations = append(operations, "r")
	}
	for i := 0; i < workload.Config.UpdatePercentage; i++ {
		operations = append(operations, "u")
	}
	for i := 0; i < workload.Config.DeletePercentage; i++ {
		operations = append(operations, "d")
	}
	for i := 0; i < workload.Config.QueryPercentage; i++ {
		operations = append(operations, "q")
	}
	if len(operations) != 100 {
		log.Fatal("Wrong workload configuration: sum of percentages is not equal 100")
	}
	for i, randI := range rand.Perm(100) {
		randOperations[i] = operations[randI]
	}
	return randOperations
}

func (w *DefaultWorkload) Something() chan string {
	operations := w.PrepareBatch()

	ch := make(chan string, 100000)

	go func() {
		for {
			ch <- operations[rand.Intn(100)]
		}
	}()

	return ch
}

// Sequentially send 100 requests
func (workload *DefaultWorkload) DoBatch(db databases.Database, state *State) {
	batch := workload.PrepareBatch()

	for _, v := range batch {
		// Increase number of passed operarions *before* batch
		// execution in order to normally share key space with
		// other workers
		if state.Operations < workload.Config.Operations {
			var err error
			state.Operations++
			switch v {
			case "c":
				state.Records++
				key := workload.GenerateNewKey(state.Records)
				value := workload.GenerateValue(key,
					workload.Config.IndexableFields, workload.Config.ValueSize)
				err = db.Create(key, value)
			case "r":
				key := workload.GenerateExistingKey(state.Records)
				err = db.Read(key)
			case "u":
				key := workload.GenerateExistingKey(state.Records)
				value := workload.GenerateValue(key,
					workload.Config.IndexableFields, workload.Config.ValueSize)
				err = db.Update(key, value)
			case "d":
				key := workload.GenerateKeyForRemoval()
				err = db.Delete(key)
			case "q":
				fieldName, fieldValue, limit := workload.GenerateQuery(
					workload.Config.IndexableFields, state.Records)
				err = db.Query(fieldName, fieldValue, limit)
			}
			if err != nil {
				state.Errors[v]++
				state.Errors["total"]++
			}
		}
	}
}

// Continuously run batches of operations
func (workload *DefaultWorkload) RunWorkload(database databases.Database,
	state *State, wg *sync.WaitGroup) {
	defer wg.Done()

	// Calculate target time for batch execution. +Inf if not defined
	targetBatchTimeF := float64(100) /
		float64(workload.Config.TargetThroughput)

	for state.Operations < workload.Config.Operations {
		// Send batch of request and measure execution time
		t0 := time.Now()
		workload.DoBatch(database, state)
		t1 := time.Now()

		// Sleep if necessary
		if !math.IsInf(targetBatchTimeF, 0) {
			targetBatchTime := time.Duration(targetBatchTimeF)
			actualBatchTime := t1.Sub(t0)
			sleepTime := (targetBatchTime - actualBatchTime)
			if sleepTime > 0 {
				time.Sleep(time.Duration(sleepTime))
			}
		}
	}
}
<commit_msg>fix sleep time<commit_after>package workloads

import (
	"bytes"
	"log"
	"math"
	"math/rand"
	"strconv"
	"sync"
	"time"

	"github.com/pavel-paulau/blurr/databases"
)

type DefaultWorkload struct {
	Config       Config
	DeletedItems int64
}

func (workload *DefaultWorkload) Init(config Config) {
	workload.Config = config
}

// Generate new *unique* key
func (workload *DefaultWorkload) GenerateNewKey(currentRecords int64) string {
	strCurrentRecords := strconv.FormatInt(currentRecords, 10)
	return Hash(strCurrentRecords)
}

// Generate random key from current key space
func (workload *DefaultWorkload) GenerateExistingKey(currentRecords int64) string {
	rand.Seed(time.Now().UnixNano())
	randRecord := workload.DeletedItems + rand.Int63n(currentRecords-workload.DeletedItems)
	strRandRecord := strconv.FormatInt(randRecord, 10)
	return Hash(strRandRecord)
}

// Generate sequential key for removal
func (workload *DefaultWorkload) GenerateKeyForRemoval() string {
	keyForRemoval := strconv.FormatInt(workload.DeletedItems+1, 10)
	workload.DeletedItems++
	return Hash(keyForRemoval)
}

// Generate value with deterministic indexable fields and arbitrary body
func (workload *DefaultWorkload) GenerateValue(key string,
	indexableFields, size int) map[string]interface{} {

	// Hex lengh is 32 characters, so only 22 indexable fields are allowed
	if indexableFields >= 20 {
		log.Fatal("Too much fields! It must be less than 20")
	}
	// Gererate indexable fields (shifting over key name)
	value := map[string]interface{}{}
	for i := 0; i < indexableFields; i++ {
		fieldName := "field" + strconv.Itoa(i)
		value[fieldName] = fieldName + "-" + key[i:i+10]
	}
	// Generate value body in order to meet value size specification
	fieldName := "field" + strconv.Itoa(indexableFields)
	buffer := bytes.Buffer{}
	bodyHash := Hash(key)
	iterations := (size - len(fieldName+"-"+key[:10])*indexableFields) / 32
	for i := 0; i < iterations; i++ {
		buffer.WriteString(bodyHash)
	}
	value[fieldName] = buffer.String()
	return value
}

func (workload *DefaultWorkload) GenerateQuery(indexableFields int,
	currentRecords int64) (string, string, int) {

	i := rand.Intn(indexableFields)
	fieldName := "field" + strconv.Itoa(i)
	fieldValue := fieldName + "-" + workload.GenerateExistingKey(currentRecords)[i:i+10]
	limit := 10 + rand.Intn(10)
	return fieldName, fieldValue, limit
}

// Generate slice of shuffled characters (CRUD-Q shorthands)
func (workload *DefaultWorkload) PrepareBatch() []string {
	operations := make([]string, 0, 100)
	randOperations := make([]string, 100, 100)
	for i := 0; i < workload.Config.CreatePercentage; i++ {
		operations = append(operations, "c")
	}
	for i := 0; i < workload.Config.ReadPercentage; i++ {
		operations = append(operations, "r")
	}
	for i := 0; i < workload.Config.UpdatePercentage; i++ {
		operations = append(operations, "u")
	}
	for i := 0; i < workload.Config.DeletePercentage; i++ {
		operations = append(operations, "d")
	}
	for i := 0; i < workload.Config.QueryPercentage; i++ {
		operations = append(operations, "q")
	}
	if len(operations) != 100 {
		log.Fatal("Wrong workload configuration: sum of percentages is not equal 100")
	}
	for i, randI := range rand.Perm(100) {
		randOperations[i] = operations[randI]
	}
	return randOperations
}

func (w *DefaultWorkload) Something() chan string {
	operations := w.PrepareBatch()

	ch := make(chan string, 100000)

	go func() {
		for {
			ch <- operations[rand.Intn(100)]
		}
	}()

	return ch
}

// Sequentially send 100 requests
func (workload *DefaultWorkload) DoBatch(db databases.Database, state *State) {
	batch := workload.PrepareBatch()

	for _, v := range batch {
		// Increase number of passed operarions *before* batch
		// execution in order to normally share key space with
		// other workers
		if state.Operations < workload.Config.Operations {
			var err error
			state.Operations++
			switch v {
			case "c":
				state.Records++
				key := workload.GenerateNewKey(state.Records)
				value := workload.GenerateValue(key,
					workload.Config.IndexableFields, workload.Config.ValueSize)
				err = db.Create(key, value)
			case "r":
				key := workload.GenerateExistingKey(state.Records)
				err = db.Read(key)
			case "u":
				key := workload.GenerateExistingKey(state.Records)
				value := workload.GenerateValue(key,
					workload.Config.IndexableFields, workload.Config.ValueSize)
				err = db.Update(key, value)
			case "d":
				key := workload.GenerateKeyForRemoval()
				err = db.Delete(key)
			case "q":
				fieldName, fieldValue, limit := workload.GenerateQuery(
					workload.Config.IndexableFields, state.Records)
				err = db.Query(fieldName, fieldValue, limit)
			}
			if err != nil {
				state.Errors[v]++
				state.Errors["total"]++
			}
		}
	}
}

// Continuously run batches of operations
func (workload *DefaultWorkload) RunWorkload(database databases.Database,
	state *State, wg *sync.WaitGroup) {
	defer wg.Done()

	// Calculate target time for batch execution. +Inf if not defined
	targetBatchTimeF := float64(100) /
		float64(workload.Config.TargetThroughput)

	for state.Operations < workload.Config.Operations {
		// Send batch of request and measure execution time
		t0 := time.Now()
		workload.DoBatch(database, state)
		t1 := time.Now()

		// Sleep if necessary
		if !math.IsInf(targetBatchTimeF, 0) {
			targetBatchTime := time.Duration(targetBatchTimeF * math.Pow10(9))
			actualBatchTime := t1.Sub(t0)
			sleepTime := (targetBatchTime - actualBatchTime)
			if sleepTime > 0 {
				time.Sleep(time.Duration(sleepTime))
			}
		}
	}
}
<|endoftext|>


<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.

// Package expvar provides a standardized interface to public variables, such
// as operation counters in servers. It exposes these variables via HTTP at
// /debug/vars in JSON format.
//
// Operations to set or modify these public variables are atomic.
//
// In addition to adding the HTTP handler, this package registers the
// following variables:
//
//	cmdline   os.Args
//	memstats  runtime.Memstats
//
// The package is sometimes only imported for the side effect of
// registering its HTTP handler and the above variables.  To use it
// this way, link this package into your program:
//	import _ "expvar"
//
package expvar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
)

// Var is an abstract type for all exported variables.
type Var interface {
	String() string
}

// Int is a 64-bit integer variable that satisfies the Var interface.
type Int struct {
	mu sync.RWMutex
	i  int64
}

func (v *Int) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return strconv.FormatInt(v.i, 10)
}

func (v *Int) Add(delta int64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.i += delta
}

func (v *Int) Set(value int64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.i = value
}

// Float is a 64-bit float variable that satisfies the Var interface.
type Float struct {
	mu sync.RWMutex
	f  float64
}

func (v *Float) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return strconv.FormatFloat(v.f, 'g', -1, 64)
}

// Add adds delta to v.
func (v *Float) Add(delta float64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.f += delta
}

// Set sets v to value.
func (v *Float) Set(value float64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.f = value
}

// Map is a string-to-Var map variable that satisfies the Var interface.
type Map struct {
	mu   sync.RWMutex
	m    map[string]Var
	keys []string // sorted
}

// KeyValue represents a single entry in a Map.
type KeyValue struct {
	Key   string
	Value Var
}

func (v *Map) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	var b bytes.Buffer
	fmt.Fprintf(&b, "{")
	first := true
	v.Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(&b, ", ")
		}
		fmt.Fprintf(&b, "\"%s\": %v", kv.Key, kv.Value)
		first = false
	})
	fmt.Fprintf(&b, "}")
	return b.String()
}

func (v *Map) Init() *Map {
	v.m = make(map[string]Var)
	return v
}

// updateKeys updates the sorted list of keys in v.keys.
// must be called with v.mu held.
func (v *Map) updateKeys() {
	if len(v.m) == len(v.keys) {
		// No new key.
		return
	}
	v.keys = v.keys[:0]
	for k := range v.m {
		v.keys = append(v.keys, k)
	}
	sort.Strings(v.keys)
}

func (v *Map) Get(key string) Var {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return v.m[key]
}

func (v *Map) Set(key string, av Var) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.m[key] = av
	v.updateKeys()
}

func (v *Map) Add(key string, delta int64) {
	v.mu.RLock()
	av, ok := v.m[key]
	v.mu.RUnlock()
	if !ok {
		// check again under the write lock
		v.mu.Lock()
		av, ok = v.m[key]
		if !ok {
			av = new(Int)
			v.m[key] = av
			v.updateKeys()
		}
		v.mu.Unlock()
	}

	// Add to Int; ignore otherwise.
	if iv, ok := av.(*Int); ok {
		iv.Add(delta)
	}
}

// AddFloat adds delta to the *Float value stored under the given map key.
func (v *Map) AddFloat(key string, delta float64) {
	v.mu.RLock()
	av, ok := v.m[key]
	v.mu.RUnlock()
	if !ok {
		// check again under the write lock
		v.mu.Lock()
		av, ok = v.m[key]
		if !ok {
			av = new(Float)
			v.m[key] = av
			v.updateKeys()
		}
		v.mu.Unlock()
	}

	// Add to Float; ignore otherwise.
	if iv, ok := av.(*Float); ok {
		iv.Add(delta)
	}
}

// Do calls f for each entry in the map.
// The map is locked during the iteration,
// but existing entries may be concurrently updated.
func (v *Map) Do(f func(KeyValue)) {
	v.mu.RLock()
	defer v.mu.RUnlock()
	for _, k := range v.keys {
		f(KeyValue{k, v.m[k]})
	}
}

// String is a string variable, and satisfies the Var interface.
type String struct {
	mu sync.RWMutex
	s  string
}

func (v *String) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return strconv.Quote(v.s)
}

func (v *String) Set(value string) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.s = value
}

// Func implements Var by calling the function
// and formatting the returned value using JSON.
type Func func() interface{}

func (f Func) String() string {
	v, _ := json.Marshal(f())
	return string(v)
}

// All published variables.
var (
	mutex   sync.RWMutex
	vars    = make(map[string]Var)
	varKeys []string // sorted
)

// Publish declares a named exported variable. This should be called from a
// package's init function when it creates its Vars. If the name is already
// registered then this will log.Panic.
func Publish(name string, v Var) {
	mutex.Lock()
	defer mutex.Unlock()
	if _, existing := vars[name]; existing {
		log.Panicln("Reuse of exported var name:", name)
	}
	vars[name] = v
	varKeys = append(varKeys, name)
	sort.Strings(varKeys)
}

// Get retrieves a named exported variable.
func Get(name string) Var {
	mutex.RLock()
	defer mutex.RUnlock()
	return vars[name]
}

// Convenience functions for creating new exported variables.

func NewInt(name string) *Int {
	v := new(Int)
	Publish(name, v)
	return v
}

func NewFloat(name string) *Float {
	v := new(Float)
	Publish(name, v)
	return v
}

func NewMap(name string) *Map {
	v := new(Map).Init()
	Publish(name, v)
	return v
}

func NewString(name string) *String {
	v := new(String)
	Publish(name, v)
	return v
}

// Do calls f for each exported variable.
// The global variable map is locked during the iteration,
// but existing entries may be concurrently updated.
func Do(f func(KeyValue)) {
	mutex.RLock()
	defer mutex.RUnlock()
	for _, k := range varKeys {
		f(KeyValue{k, vars[k]})
	}
}

func expvarHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	fmt.Fprintf(w, "{\n")
	first := true
	Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(w, ",\n")
		}
		first = false
		fmt.Fprintf(w, "%q: %s", kv.Key, kv.Value)
	})
	fmt.Fprintf(w, "\n}\n")
}

func cmdline() interface{} {
	return os.Args
}

func memstats() interface{} {
	stats := new(runtime.MemStats)
	runtime.ReadMemStats(stats)
	return *stats
}

func init() {
	http.HandleFunc("/debug/vars", expvarHandler)
	Publish("cmdline", Func(cmdline))
	Publish("memstats", Func(memstats))
}
<commit_msg>expvar: don't recursively acquire Map.RLock<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.

// Package expvar provides a standardized interface to public variables, such
// as operation counters in servers. It exposes these variables via HTTP at
// /debug/vars in JSON format.
//
// Operations to set or modify these public variables are atomic.
//
// In addition to adding the HTTP handler, this package registers the
// following variables:
//
//	cmdline   os.Args
//	memstats  runtime.Memstats
//
// The package is sometimes only imported for the side effect of
// registering its HTTP handler and the above variables.  To use it
// this way, link this package into your program:
//	import _ "expvar"
//
package expvar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
)

// Var is an abstract type for all exported variables.
type Var interface {
	String() string
}

// Int is a 64-bit integer variable that satisfies the Var interface.
type Int struct {
	mu sync.RWMutex
	i  int64
}

func (v *Int) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return strconv.FormatInt(v.i, 10)
}

func (v *Int) Add(delta int64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.i += delta
}

func (v *Int) Set(value int64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.i = value
}

// Float is a 64-bit float variable that satisfies the Var interface.
type Float struct {
	mu sync.RWMutex
	f  float64
}

func (v *Float) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return strconv.FormatFloat(v.f, 'g', -1, 64)
}

// Add adds delta to v.
func (v *Float) Add(delta float64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.f += delta
}

// Set sets v to value.
func (v *Float) Set(value float64) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.f = value
}

// Map is a string-to-Var map variable that satisfies the Var interface.
type Map struct {
	mu   sync.RWMutex
	m    map[string]Var
	keys []string // sorted
}

// KeyValue represents a single entry in a Map.
type KeyValue struct {
	Key   string
	Value Var
}

func (v *Map) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	var b bytes.Buffer
	fmt.Fprintf(&b, "{")
	first := true
	v.doLocked(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(&b, ", ")
		}
		fmt.Fprintf(&b, "\"%s\": %v", kv.Key, kv.Value)
		first = false
	})
	fmt.Fprintf(&b, "}")
	return b.String()
}

func (v *Map) Init() *Map {
	v.m = make(map[string]Var)
	return v
}

// updateKeys updates the sorted list of keys in v.keys.
// must be called with v.mu held.
func (v *Map) updateKeys() {
	if len(v.m) == len(v.keys) {
		// No new key.
		return
	}
	v.keys = v.keys[:0]
	for k := range v.m {
		v.keys = append(v.keys, k)
	}
	sort.Strings(v.keys)
}

func (v *Map) Get(key string) Var {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return v.m[key]
}

func (v *Map) Set(key string, av Var) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.m[key] = av
	v.updateKeys()
}

func (v *Map) Add(key string, delta int64) {
	v.mu.RLock()
	av, ok := v.m[key]
	v.mu.RUnlock()
	if !ok {
		// check again under the write lock
		v.mu.Lock()
		av, ok = v.m[key]
		if !ok {
			av = new(Int)
			v.m[key] = av
			v.updateKeys()
		}
		v.mu.Unlock()
	}

	// Add to Int; ignore otherwise.
	if iv, ok := av.(*Int); ok {
		iv.Add(delta)
	}
}

// AddFloat adds delta to the *Float value stored under the given map key.
func (v *Map) AddFloat(key string, delta float64) {
	v.mu.RLock()
	av, ok := v.m[key]
	v.mu.RUnlock()
	if !ok {
		// check again under the write lock
		v.mu.Lock()
		av, ok = v.m[key]
		if !ok {
			av = new(Float)
			v.m[key] = av
			v.updateKeys()
		}
		v.mu.Unlock()
	}

	// Add to Float; ignore otherwise.
	if iv, ok := av.(*Float); ok {
		iv.Add(delta)
	}
}

// Do calls f for each entry in the map.
// The map is locked during the iteration,
// but existing entries may be concurrently updated.
func (v *Map) Do(f func(KeyValue)) {
	v.mu.RLock()
	defer v.mu.RUnlock()
	v.doLocked(f)
}

// doRLocked calls f for each entry in the map.
// v.mu must be held for reads.
func (v *Map) doLocked(f func(KeyValue)) {
	for _, k := range v.keys {
		f(KeyValue{k, v.m[k]})
	}
}

// String is a string variable, and satisfies the Var interface.
type String struct {
	mu sync.RWMutex
	s  string
}

func (v *String) String() string {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return strconv.Quote(v.s)
}

func (v *String) Set(value string) {
	v.mu.Lock()
	defer v.mu.Unlock()
	v.s = value
}

// Func implements Var by calling the function
// and formatting the returned value using JSON.
type Func func() interface{}

func (f Func) String() string {
	v, _ := json.Marshal(f())
	return string(v)
}

// All published variables.
var (
	mutex   sync.RWMutex
	vars    = make(map[string]Var)
	varKeys []string // sorted
)

// Publish declares a named exported variable. This should be called from a
// package's init function when it creates its Vars. If the name is already
// registered then this will log.Panic.
func Publish(name string, v Var) {
	mutex.Lock()
	defer mutex.Unlock()
	if _, existing := vars[name]; existing {
		log.Panicln("Reuse of exported var name:", name)
	}
	vars[name] = v
	varKeys = append(varKeys, name)
	sort.Strings(varKeys)
}

// Get retrieves a named exported variable.
func Get(name string) Var {
	mutex.RLock()
	defer mutex.RUnlock()
	return vars[name]
}

// Convenience functions for creating new exported variables.

func NewInt(name string) *Int {
	v := new(Int)
	Publish(name, v)
	return v
}

func NewFloat(name string) *Float {
	v := new(Float)
	Publish(name, v)
	return v
}

func NewMap(name string) *Map {
	v := new(Map).Init()
	Publish(name, v)
	return v
}

func NewString(name string) *String {
	v := new(String)
	Publish(name, v)
	return v
}

// Do calls f for each exported variable.
// The global variable map is locked during the iteration,
// but existing entries may be concurrently updated.
func Do(f func(KeyValue)) {
	mutex.RLock()
	defer mutex.RUnlock()
	for _, k := range varKeys {
		f(KeyValue{k, vars[k]})
	}
}

func expvarHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	fmt.Fprintf(w, "{\n")
	first := true
	Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(w, ",\n")
		}
		first = false
		fmt.Fprintf(w, "%q: %s", kv.Key, kv.Value)
	})
	fmt.Fprintf(w, "\n}\n")
}

func cmdline() interface{} {
	return os.Args
}

func memstats() interface{} {
	stats := new(runtime.MemStats)
	runtime.ReadMemStats(stats)
	return *stats
}

func init() {
	http.HandleFunc("/debug/vars", expvarHandler)
	Publish("cmdline", Func(cmdline))
	Publish("memstats", Func(memstats))
}
<|endoftext|>


<commit_before>// Copyright 2010 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.

// Package dict implements the Dictionary Server Protocol
// as defined in RFC 2229.
package dict

import (
	"container/vector"
	"net/textproto"
	"os"
	"strconv"
	"strings"
)

// A Client represents a client connection to a dictionary server.
type Client struct {
	text *textproto.Conn
}

// Dial returns a new client connected to a dictionary server at
// addr on the given network.
func Dial(network, addr string) (*Client, os.Error) {
	text, err := textproto.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	_, _, err = text.ReadCodeLine(220)
	if err != nil {
		text.Close()
		return nil, err
	}
	return &Client{text: text}, nil
}

// Close closes the connection to the dictionary server.
func (c *Client) Close() os.Error {
	return c.text.Close()
}

// A Dict represents a dictionary available on the server.
type Dict struct {
	Name string // short name of dictionary
	Desc string // long description
}

// Dicts returns a list of the dictionaries available on the server.
func (c *Client) Dicts() ([]Dict, os.Error) {
	id, err := c.text.Cmd("SHOW DB")
	if err != nil {
		return nil, err
	}

	c.text.StartResponse(id)
	defer c.text.EndResponse(id)

	lines, err := c.text.ReadDotLines()
	if err != nil {
		return nil, err
	}
	_, _, err = c.text.ReadCodeLine(250)

	dicts := make([]Dict, len(lines))
	for i := range dicts {
		d := &dicts[i]
		a, _ := fields(lines[i])
		if len(a) < 2 {
			return nil, textproto.ProtocolError("invalid dictionary: " + lines[i])
		}
		d.Name = a[0]
		d.Desc = a[1]
	}
	return dicts, err
}

// A Defn represents a definition.
type Defn struct {
	Dict Dict   // Dict where definition was found
	Word string // Word being defined
	Text []byte // Definition text, typically multiple lines
}

// Define requests the definition of the given word.
// The argument dict names the dictionary to use,
// the Name field of a Dict returned by Dicts.
//
// The special dictionary name "!" means to look in all the
// server's dictionaries.
// The special dictionary name "*" means to look in all the
// server's dictionaries in turn, stopping after finding the word
// in one of them.
func (c *Client) Define(dict, word string) ([]*Defn, os.Error) {
	id, err := c.text.Cmd("DEFINE %s %q", dict, word)
	if err != nil {
		return nil, err
	}

	c.text.StartResponse(id)
	defer c.text.EndResponse(id)

	_, line, err := c.text.ReadCodeLine(150)
	a, _ := fields(line)
	if len(a) < 1 {
		return nil, textproto.ProtocolError("malformed response: " + line)
	}
	n, err := strconv.Atoi(a[0])
	if err != nil {
		return nil, textproto.ProtocolError("invalid definition count: " + a[0])
	}
	def := make([]*Defn, n)
	for i := 0; i < n; i++ {
		_, line, err = c.text.ReadCodeLine(151)
		if err != nil {
			return nil, err
		}
		a, _ := fields(line)
		if len(a) < 3 {
			// skip it, to keep protocol in sync
			i--
			n--
			def = def[0:n]
			continue
		}
		d := &Defn{Word: a[0], Dict: Dict{a[1], a[2]}}
		d.Text, err = c.text.ReadDotBytes()
		if err != nil {
			return nil, err
		}
		def[i] = d
	}
	_, _, err = c.text.ReadCodeLine(250)
	return def, err
}

// Fields returns the fields in s.
// Fields are space separated unquoted words
// or quoted with single or double quote.
func fields(s string) ([]string, os.Error) {
	var v vector.StringVector
	i := 0
	for {
		for i < len(s) && (s[i] == ' ' || s[i] == '\t') {
			i++
		}
		if i >= len(s) {
			break
		}
		if s[i] == '"' || s[i] == '\'' {
			q := s[i]
			// quoted string
			var j int
			for j = i + 1; ; j++ {
				if j >= len(s) {
					return nil, textproto.ProtocolError("malformed quoted string")
				}
				if s[j] == '\\' {
					j++
					continue
				}
				if s[j] == q {
					j++
					break
				}
			}
			v.Push(unquote(s[i+1 : j-1]))
			i = j
		} else {
			// atom
			var j int
			for j = i; j < len(s); j++ {
				if s[j] == ' ' || s[j] == '\t' || s[j] == '\\' || s[j] == '"' || s[j] == '\'' {
					break
				}
			}
			v.Push(s[i:j])
			i = j
		}
		if i < len(s) {
			c := s[i]
			if c != ' ' && c != '\t' {
				return nil, textproto.ProtocolError("quotes not on word boundaries")
			}
		}
	}
	return v, nil
}

func unquote(s string) string {
	if strings.Index(s, "\\") < 0 {
		return s
	}
	b := []byte(s)
	w := 0
	for r := 0; r < len(b); r++ {
		c := b[r]
		if c == '\\' {
			r++
			c = b[r]
		}
		b[w] = c
		w++
	}
	return string(b[0:w])
}
<commit_msg>net/dict: parse response correctly + typo<commit_after>// Copyright 2010 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.

// Package dict implements the Dictionary Server Protocol
// as defined in RFC 2229.
package dict

import (
	"container/vector"
	"net/textproto"
	"os"
	"strconv"
	"strings"
)

// A Client represents a client connection to a dictionary server.
type Client struct {
	text *textproto.Conn
}

// Dial returns a new client connected to a dictionary server at
// addr on the given network.
func Dial(network, addr string) (*Client, os.Error) {
	text, err := textproto.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	_, _, err = text.ReadCodeLine(220)
	if err != nil {
		text.Close()
		return nil, err
	}
	return &Client{text: text}, nil
}

// Close closes the connection to the dictionary server.
func (c *Client) Close() os.Error {
	return c.text.Close()
}

// A Dict represents a dictionary available on the server.
type Dict struct {
	Name string // short name of dictionary
	Desc string // long description
}

// Dicts returns a list of the dictionaries available on the server.
func (c *Client) Dicts() ([]Dict, os.Error) {
	id, err := c.text.Cmd("SHOW DB")
	if err != nil {
		return nil, err
	}

	c.text.StartResponse(id)
	defer c.text.EndResponse(id)

	_, _, err = c.text.ReadCodeLine(110)
	if err != nil {
		return nil, err
	}
	lines, err := c.text.ReadDotLines()
	if err != nil {
		return nil, err
	}
	_, _, err = c.text.ReadCodeLine(250)

	dicts := make([]Dict, len(lines))
	for i := range dicts {
		d := &dicts[i]
		a, _ := fields(lines[i])
		if len(a) < 2 {
			return nil, textproto.ProtocolError("invalid dictionary: " + lines[i])
		}
		d.Name = a[0]
		d.Desc = a[1]
	}
	return dicts, err
}

// A Defn represents a definition.
type Defn struct {
	Dict Dict   // Dict where definition was found
	Word string // Word being defined
	Text []byte // Definition text, typically multiple lines
}

// Define requests the definition of the given word.
// The argument dict names the dictionary to use,
// the Name field of a Dict returned by Dicts.
//
// The special dictionary name "*" means to look in all the
// server's dictionaries.
// The special dictionary name "!" means to look in all the
// server's dictionaries in turn, stopping after finding the word
// in one of them.
func (c *Client) Define(dict, word string) ([]*Defn, os.Error) {
	id, err := c.text.Cmd("DEFINE %s %q", dict, word)
	if err != nil {
		return nil, err
	}

	c.text.StartResponse(id)
	defer c.text.EndResponse(id)

	_, line, err := c.text.ReadCodeLine(150)
	if err != nil {
		return nil, err
	}
	a, _ := fields(line)
	if len(a) < 1 {
		return nil, textproto.ProtocolError("malformed response: " + line)
	}
	n, err := strconv.Atoi(a[0])
	if err != nil {
		return nil, textproto.ProtocolError("invalid definition count: " + a[0])
	}
	def := make([]*Defn, n)
	for i := 0; i < n; i++ {
		_, line, err = c.text.ReadCodeLine(151)
		if err != nil {
			return nil, err
		}
		a, _ := fields(line)
		if len(a) < 3 {
			// skip it, to keep protocol in sync
			i--
			n--
			def = def[0:n]
			continue
		}
		d := &Defn{Word: a[0], Dict: Dict{a[1], a[2]}}
		d.Text, err = c.text.ReadDotBytes()
		if err != nil {
			return nil, err
		}
		def[i] = d
	}
	_, _, err = c.text.ReadCodeLine(250)
	return def, err
}

// Fields returns the fields in s.
// Fields are space separated unquoted words
// or quoted with single or double quote.
func fields(s string) ([]string, os.Error) {
	var v vector.StringVector
	i := 0
	for {
		for i < len(s) && (s[i] == ' ' || s[i] == '\t') {
			i++
		}
		if i >= len(s) {
			break
		}
		if s[i] == '"' || s[i] == '\'' {
			q := s[i]
			// quoted string
			var j int
			for j = i + 1; ; j++ {
				if j >= len(s) {
					return nil, textproto.ProtocolError("malformed quoted string")
				}
				if s[j] == '\\' {
					j++
					continue
				}
				if s[j] == q {
					j++
					break
				}
			}
			v.Push(unquote(s[i+1 : j-1]))
			i = j
		} else {
			// atom
			var j int
			for j = i; j < len(s); j++ {
				if s[j] == ' ' || s[j] == '\t' || s[j] == '\\' || s[j] == '"' || s[j] == '\'' {
					break
				}
			}
			v.Push(s[i:j])
			i = j
		}
		if i < len(s) {
			c := s[i]
			if c != ' ' && c != '\t' {
				return nil, textproto.ProtocolError("quotes not on word boundaries")
			}
		}
	}
	return v, nil
}

func unquote(s string) string {
	if strings.Index(s, "\\") < 0 {
		return s
	}
	b := []byte(s)
	w := 0
	for r := 0; r < len(b); r++ {
		c := b[r]
		if c == '\\' {
			r++
			c = b[r]
		}
		b[w] = c
		w++
	}
	return string(b[0:w])
}
<|endoftext|>


<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.

// +build darwin freebsd linux netbsd openbsd windows

package os

import (
	"syscall"
)

// StartProcess starts a new process with the program, arguments and attributes
// specified by name, argv and attr.
//
// StartProcess is a low-level interface. The os/exec package provides
// higher-level interfaces.
//
// If there is an error, it will be of type *PathError.
func StartProcess(name string, argv []string, attr *ProcAttr) (p *Process, err error) {
	sysattr := &syscall.ProcAttr{
		Dir: attr.Dir,
		Env: attr.Env,
		Sys: attr.Sys,
	}
	if sysattr.Env == nil {
		sysattr.Env = Environ()
	}
	for _, f := range attr.Files {
		sysattr.Files = append(sysattr.Files, f.Fd())
	}

	pid, h, e := syscall.StartProcess(name, argv, sysattr)
	if e != nil {
		return nil, &PathError{"fork/exec", name, e}
	}
	return newProcess(pid, h), nil
}

// Kill causes the Process to exit immediately.
func (p *Process) Kill() error {
	return p.Signal(Kill)
}

// ProcessState stores information about process as reported by Wait.
type ProcessState struct {
	pid    int                // The process's id.
	status syscall.WaitStatus // System-dependent status info.
	rusage *syscall.Rusage
}

// Pid returns the process id of the exited process.
func (p *ProcessState) Pid() int {
	return p.pid
}

// Exited returns whether the program has exited.
func (p *ProcessState) Exited() bool {
	return p.status.Exited()
}

// Success reports whether the program exited successfully,
// such as with exit status 0 on Unix.
func (p *ProcessState) Success() bool {
	return p.status.ExitStatus() == 0
}

// Sys returns system-dependent exit information about
// the process.  Convert it to the appropriate underlying
// type, such as syscall.WaitStatus on Unix, to access its contents.
func (p *ProcessState) Sys() interface{} {
	return p.status
}

// SysUsage returns system-dependent resource usage information about
// the exited process.  Convert it to the appropriate underlying
// type, such as *syscall.Rusage on Unix, to access its contents.
func (p *ProcessState) SysUsage() interface{} {
	return p.rusage
}

// Convert i to decimal string.
func itod(i int) string {
	if i == 0 {
		return "0"
	}

	u := uint64(i)
	if i < 0 {
		u = -u
	}

	// Assemble decimal in reverse order.
	var b [32]byte
	bp := len(b)
	for ; u > 0; u /= 10 {
		bp--
		b[bp] = byte(u%10) + '0'
	}

	if i < 0 {
		bp--
		b[bp] = '-'
	}

	return string(b[bp:])
}

func (p *ProcessState) String() string {
	if p == nil {
		return "<nil>"
	}
	status := p.Sys().(syscall.WaitStatus)
	res := ""
	switch {
	case status.Exited():
		res = "exit status " + itod(status.ExitStatus())
	case status.Signaled():
		res = "signal " + itod(int(status.Signal()))
	case status.Stopped():
		res = "stop signal " + itod(int(status.StopSignal()))
		if status.StopSignal() == syscall.SIGTRAP && status.TrapCause() != 0 {
			res += " (trap " + itod(status.TrapCause()) + ")"
		}
	case status.Continued():
		res = "continued"
	}
	if status.CoreDump() {
		res += " (core dumped)"
	}
	return res
}
<commit_msg>os: diagnose chdir error during StartProcess<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.

// +build darwin freebsd linux netbsd openbsd windows

package os

import (
	"syscall"
)

// StartProcess starts a new process with the program, arguments and attributes
// specified by name, argv and attr.
//
// StartProcess is a low-level interface. The os/exec package provides
// higher-level interfaces.
//
// If there is an error, it will be of type *PathError.
func StartProcess(name string, argv []string, attr *ProcAttr) (p *Process, err error) {
	// Double-check existence of the directory we want
	// to chdir into.  We can make the error clearer this way.
	if attr != nil && attr.Dir != "" {
		if _, err := Stat(attr.Dir); err != nil {
			pe := err.(*PathError)
			pe.Op = "chdir"
			return nil, pe
		}
	}

	sysattr := &syscall.ProcAttr{
		Dir: attr.Dir,
		Env: attr.Env,
		Sys: attr.Sys,
	}
	if sysattr.Env == nil {
		sysattr.Env = Environ()
	}
	for _, f := range attr.Files {
		sysattr.Files = append(sysattr.Files, f.Fd())
	}

	pid, h, e := syscall.StartProcess(name, argv, sysattr)
	if e != nil {
		return nil, &PathError{"fork/exec", name, e}
	}
	return newProcess(pid, h), nil
}

// Kill causes the Process to exit immediately.
func (p *Process) Kill() error {
	return p.Signal(Kill)
}

// ProcessState stores information about process as reported by Wait.
type ProcessState struct {
	pid    int                // The process's id.
	status syscall.WaitStatus // System-dependent status info.
	rusage *syscall.Rusage
}

// Pid returns the process id of the exited process.
func (p *ProcessState) Pid() int {
	return p.pid
}

// Exited returns whether the program has exited.
func (p *ProcessState) Exited() bool {
	return p.status.Exited()
}

// Success reports whether the program exited successfully,
// such as with exit status 0 on Unix.
func (p *ProcessState) Success() bool {
	return p.status.ExitStatus() == 0
}

// Sys returns system-dependent exit information about
// the process.  Convert it to the appropriate underlying
// type, such as syscall.WaitStatus on Unix, to access its contents.
func (p *ProcessState) Sys() interface{} {
	return p.status
}

// SysUsage returns system-dependent resource usage information about
// the exited process.  Convert it to the appropriate underlying
// type, such as *syscall.Rusage on Unix, to access its contents.
func (p *ProcessState) SysUsage() interface{} {
	return p.rusage
}

// Convert i to decimal string.
func itod(i int) string {
	if i == 0 {
		return "0"
	}

	u := uint64(i)
	if i < 0 {
		u = -u
	}

	// Assemble decimal in reverse order.
	var b [32]byte
	bp := len(b)
	for ; u > 0; u /= 10 {
		bp--
		b[bp] = byte(u%10) + '0'
	}

	if i < 0 {
		bp--
		b[bp] = '-'
	}

	return string(b[bp:])
}

func (p *ProcessState) String() string {
	if p == nil {
		return "<nil>"
	}
	status := p.Sys().(syscall.WaitStatus)
	res := ""
	switch {
	case status.Exited():
		res = "exit status " + itod(status.ExitStatus())
	case status.Signaled():
		res = "signal " + itod(int(status.Signal()))
	case status.Stopped():
		res = "stop signal " + itod(int(status.StopSignal()))
		if status.StopSignal() == syscall.SIGTRAP && status.TrapCause() != 0 {
			res += " (trap " + itod(status.TrapCause()) + ")"
		}
	case status.Continued():
		res = "continued"
	}
	if status.CoreDump() {
		res += " (core dumped)"
	}
	return res
}
<|endoftext|>