<commit_before>package proto

import (
	"errors"
	"net"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

// extendDeadline is a helper function for extending the connection timeout.
func extendDeadline(conn net.Conn, d time.Duration) { _ = conn.SetDeadline(time.Now().Add(d)) }

// startRevision is run at the beginning of each revision iteration. It reads
// the host's settings confirms that the values are acceptable, and writes an acceptance.
func startRevision(conn net.Conn, host modules.HostDBEntry) error {
	// verify the host's settings and confirm its identity
	recvSettings, err := verifySettings(conn, host)
	if err != nil {
		return err
	} else if !recvSettings.AcceptingContracts {
		// no need to reject; host will already have disconnected at this point
		return errors.New("host is not accepting contracts")
	}
	return modules.WriteNegotiationAcceptance(conn)
}

// startDownload is run at the beginning of each download iteration. It reads
// the host's settings confirms that the values are acceptable, and writes an acceptance.
func startDownload(conn net.Conn, host modules.HostDBEntry) error {
	// verify the host's settings and confirm its identity
	_, err := verifySettings(conn, host)
	if err != nil {
		return err
	}
	return modules.WriteNegotiationAcceptance(conn)
}

// verifySettings reads a signed HostSettings object from conn, validates the
// signature, and checks for discrepancies between the known settings and the
// received settings. If there is a discrepancy, the hostDB is notified. The
// received settings are returned.
func verifySettings(conn net.Conn, host modules.HostDBEntry) (modules.HostDBEntry, error) {
	// convert host key (types.SiaPublicKey) to a crypto.PublicKey
	if host.PublicKey.Algorithm != types.SignatureEd25519 || len(host.PublicKey.Key) != crypto.PublicKeySize {
		build.Critical("hostdb did not filter out host with wrong signature algorithm:", host.PublicKey.Algorithm)
		return modules.HostDBEntry{}, errors.New("host used unsupported signature algorithm")
	}
	var pk crypto.PublicKey
	copy(pk[:], host.PublicKey.Key)

	// read signed host settings
	var recvSettings modules.HostExternalSettings
	if err := crypto.ReadSignedObject(conn, &recvSettings, modules.NegotiateMaxHostExternalSettingsLen, pk); err != nil {
		return modules.HostDBEntry{}, errors.New("couldn't read host's settings: " + err.Error())
	}
	// TODO: check recvSettings against host.HostExternalSettings. If there is
	// a discrepancy, write the error to conn.
	if recvSettings.NetAddress != host.NetAddress {
		// for now, just overwrite the NetAddress, since we know that
		// host.NetAddress works (it was the one we dialed to get conn)
		recvSettings.NetAddress = host.NetAddress
	}
	host.HostExternalSettings = recvSettings
	return host, nil
}

// verifyRecentRevision confirms that the host and contractor agree upon the current
// state of the contract being revised.
func verifyRecentRevision(conn net.Conn, contract modules.RenterContract) error {
	// send contract ID
	if err := encoding.WriteObject(conn, contract.ID); err != nil {
		return errors.New("couldn't send contract ID: " + err.Error())
	}
	// read challenge
	var challenge crypto.Hash
	if err := encoding.ReadObject(conn, &challenge, 32); err != nil {
		return errors.New("couldn't read challenge: " + err.Error())
	}
	// sign and return
	sig, err := crypto.SignHash(challenge, contract.SecretKey)
	if err != nil {
		return err
	} else if err := encoding.WriteObject(conn, sig); err != nil {
		return errors.New("couldn't send challenge response: " + err.Error())
	}
	// read acceptance
	if err := modules.ReadNegotiationAcceptance(conn); err != nil {
		return errors.New("host did not accept revision request: " + err.Error())
	}
	// read last revision and signatures
	var lastRevision types.FileContractRevision
	var hostSignatures []types.TransactionSignature
	if err := encoding.ReadObject(conn, &lastRevision, 2048); err != nil {
		return errors.New("couldn't read last revision: " + err.Error())
	}
	if err := encoding.ReadObject(conn, &hostSignatures, 2048); err != nil {
		return errors.New("couldn't read host signatures: " + err.Error())
	}
	// Check that the unlock hashes match; if they do not, something is
	// seriously wrong. Otherwise, check that the revision numbers match.
	if lastRevision.UnlockConditions.UnlockHash() != contract.LastRevision.UnlockConditions.UnlockHash() {
		return errors.New("unlock conditions do not match")
	} else if lastRevision.NewRevisionNumber != contract.LastRevision.NewRevisionNumber {
		return &recentRevisionError{contract.LastRevision.NewRevisionNumber, lastRevision.NewRevisionNumber}
	}
	// NOTE: we can fake the blockheight here because it doesn't affect
	// verification; it just needs to be above the fork height and below the
	// contract expiration (which was checked earlier).
	return modules.VerifyFileContractRevisionTransactionSignatures(lastRevision, hostSignatures, contract.FileContract.WindowStart-1)
}

// negotiateRevision sends a revision and actions to the host for approval,
// completing one iteration of the revision loop.
func negotiateRevision(conn net.Conn, rev types.FileContractRevision, secretKey crypto.SecretKey) (types.Transaction, error) {
	// create transaction containing the revision
	signedTxn := types.Transaction{
		FileContractRevisions: []types.FileContractRevision{rev},
		TransactionSignatures: []types.TransactionSignature{{
			ParentID:       crypto.Hash(rev.ParentID),
			CoveredFields:  types.CoveredFields{FileContractRevisions: []uint64{0}},
			PublicKeyIndex: 0, // renter key is always first -- see formContract
		}},
	}
	// sign the transaction
	encodedSig, _ := crypto.SignHash(signedTxn.SigHash(0), secretKey) // no error possible
	signedTxn.TransactionSignatures[0].Signature = encodedSig[:]

	// send the revision
	if err := encoding.WriteObject(conn, rev); err != nil {
		return types.Transaction{}, errors.New("couldn't send revision: " + err.Error())
	}
	// read acceptance
	if err := modules.ReadNegotiationAcceptance(conn); err != nil {
		return types.Transaction{}, errors.New("host did not accept revision: " + err.Error())
	}

	// send the new transaction signature
	if err := encoding.WriteObject(conn, signedTxn.TransactionSignatures[0]); err != nil {
		return types.Transaction{}, errors.New("couldn't send transaction signature: " + err.Error())
	}
	// read the host's acceptance and transaction signature
	// NOTE: if the host sends ErrStopResponse, we should continue processing
	// the revision, but return the error anyway.
	responseErr := modules.ReadNegotiationAcceptance(conn)
	if responseErr != nil && responseErr != modules.ErrStopResponse {
		return types.Transaction{}, errors.New("host did not accept transaction signature: " + responseErr.Error())
	}
	var hostSig types.TransactionSignature
	if err := encoding.ReadObject(conn, &hostSig, 16e3); err != nil {
		return types.Transaction{}, errors.New("couldn't read host's signature: " + err.Error())
	}

	// add the signature to the transaction and verify it
	// NOTE: we can fake the blockheight here because it doesn't affect
	// verification; it just needs to be above the fork height and below the
	// contract expiration (which was checked earlier).
	verificationHeight := rev.NewWindowStart - 1
	signedTxn.TransactionSignatures = append(signedTxn.TransactionSignatures, hostSig)
	if err := signedTxn.StandaloneValid(verificationHeight); err != nil {
		return types.Transaction{}, err
	}

	// if the host sent ErrStopResponse, return it
	return signedTxn, responseErr
}

// newRevision creates a copy of current with its revision number incremented,
// and with cost transferred from the renter to the host.
func newRevision(current types.FileContractRevision, cost types.Currency) types.FileContractRevision {
	rev := current

	// need to manually copy slice memory
	rev.NewValidProofOutputs = make([]types.SiacoinOutput, 2)
	rev.NewMissedProofOutputs = make([]types.SiacoinOutput, 3)
	copy(rev.NewValidProofOutputs, current.NewValidProofOutputs)
	copy(rev.NewMissedProofOutputs, current.NewMissedProofOutputs)

	// move valid payout from renter to host
	rev.NewValidProofOutputs[0].Value = current.NewValidProofOutputs[0].Value.Sub(cost)
	rev.NewValidProofOutputs[1].Value = current.NewValidProofOutputs[1].Value.Add(cost)

	// move missed payout from renter to void
	rev.NewMissedProofOutputs[0].Value = current.NewMissedProofOutputs[0].Value.Sub(cost)
	rev.NewMissedProofOutputs[2].Value = current.NewMissedProofOutputs[2].Value.Add(cost)

	// increment revision number
	rev.NewRevisionNumber++

	return rev
}

// newDownloadRevision revises the current revision to cover the cost of
// downloading data.
func newDownloadRevision(current types.FileContractRevision, downloadCost types.Currency) types.FileContractRevision {
	return newRevision(current, downloadCost)
}

// newUploadRevision revises the current revision to cover the cost of
// uploading a sector.
func newUploadRevision(current types.FileContractRevision, merkleRoot crypto.Hash, price, collateral types.Currency) types.FileContractRevision {
	rev := newRevision(current, price)

	// move collateral from host to void
	rev.NewMissedProofOutputs[1].Value = rev.NewMissedProofOutputs[1].Value.Sub(collateral)
	rev.NewMissedProofOutputs[2].Value = rev.NewMissedProofOutputs[2].Value.Add(collateral)

	// set new filesize and Merkle root
	rev.NewFileSize += modules.SectorSize
	rev.NewFileMerkleRoot = merkleRoot
	return rev
}

// newDeleteRevision revises the current revision to cover the cost of
// deleting a sector.
func newDeleteRevision(current types.FileContractRevision, merkleRoot crypto.Hash) types.FileContractRevision {
	rev := newRevision(current, types.ZeroCurrency)
	rev.NewFileSize -= modules.SectorSize
	rev.NewFileMerkleRoot = merkleRoot
	return rev
}

// newModifyRevision revises the current revision to cover the cost of
// modifying a sector.
func newModifyRevision(current types.FileContractRevision, merkleRoot crypto.Hash, uploadCost types.Currency) types.FileContractRevision {
	rev := newRevision(current, uploadCost)
	rev.NewFileMerkleRoot = merkleRoot
	return rev
}
<commit_msg>don't check AcceptingContracts during revision (#1375)<commit_after>package proto

import (
	"errors"
	"net"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

// extendDeadline is a helper function for extending the connection timeout.
func extendDeadline(conn net.Conn, d time.Duration) { _ = conn.SetDeadline(time.Now().Add(d)) }

// startRevision is run at the beginning of each revision iteration. It reads
// the host's settings confirms that the values are acceptable, and writes an acceptance.
func startRevision(conn net.Conn, host modules.HostDBEntry) error {
	// verify the host's settings and confirm its identity
	_, err := verifySettings(conn, host)
	if err != nil {
		return err
	}
	return modules.WriteNegotiationAcceptance(conn)
}

// startDownload is run at the beginning of each download iteration. It reads
// the host's settings confirms that the values are acceptable, and writes an acceptance.
func startDownload(conn net.Conn, host modules.HostDBEntry) error {
	// verify the host's settings and confirm its identity
	_, err := verifySettings(conn, host)
	if err != nil {
		return err
	}
	return modules.WriteNegotiationAcceptance(conn)
}

// verifySettings reads a signed HostSettings object from conn, validates the
// signature, and checks for discrepancies between the known settings and the
// received settings. If there is a discrepancy, the hostDB is notified. The
// received settings are returned.
func verifySettings(conn net.Conn, host modules.HostDBEntry) (modules.HostDBEntry, error) {
	// convert host key (types.SiaPublicKey) to a crypto.PublicKey
	if host.PublicKey.Algorithm != types.SignatureEd25519 || len(host.PublicKey.Key) != crypto.PublicKeySize {
		build.Critical("hostdb did not filter out host with wrong signature algorithm:", host.PublicKey.Algorithm)
		return modules.HostDBEntry{}, errors.New("host used unsupported signature algorithm")
	}
	var pk crypto.PublicKey
	copy(pk[:], host.PublicKey.Key)

	// read signed host settings
	var recvSettings modules.HostExternalSettings
	if err := crypto.ReadSignedObject(conn, &recvSettings, modules.NegotiateMaxHostExternalSettingsLen, pk); err != nil {
		return modules.HostDBEntry{}, errors.New("couldn't read host's settings: " + err.Error())
	}
	// TODO: check recvSettings against host.HostExternalSettings. If there is
	// a discrepancy, write the error to conn.
	if recvSettings.NetAddress != host.NetAddress {
		// for now, just overwrite the NetAddress, since we know that
		// host.NetAddress works (it was the one we dialed to get conn)
		recvSettings.NetAddress = host.NetAddress
	}
	host.HostExternalSettings = recvSettings
	return host, nil
}

// verifyRecentRevision confirms that the host and contractor agree upon the current
// state of the contract being revised.
func verifyRecentRevision(conn net.Conn, contract modules.RenterContract) error {
	// send contract ID
	if err := encoding.WriteObject(conn, contract.ID); err != nil {
		return errors.New("couldn't send contract ID: " + err.Error())
	}
	// read challenge
	var challenge crypto.Hash
	if err := encoding.ReadObject(conn, &challenge, 32); err != nil {
		return errors.New("couldn't read challenge: " + err.Error())
	}
	// sign and return
	sig, err := crypto.SignHash(challenge, contract.SecretKey)
	if err != nil {
		return err
	} else if err := encoding.WriteObject(conn, sig); err != nil {
		return errors.New("couldn't send challenge response: " + err.Error())
	}
	// read acceptance
	if err := modules.ReadNegotiationAcceptance(conn); err != nil {
		return errors.New("host did not accept revision request: " + err.Error())
	}
	// read last revision and signatures
	var lastRevision types.FileContractRevision
	var hostSignatures []types.TransactionSignature
	if err := encoding.ReadObject(conn, &lastRevision, 2048); err != nil {
		return errors.New("couldn't read last revision: " + err.Error())
	}
	if err := encoding.ReadObject(conn, &hostSignatures, 2048); err != nil {
		return errors.New("couldn't read host signatures: " + err.Error())
	}
	// Check that the unlock hashes match; if they do not, something is
	// seriously wrong. Otherwise, check that the revision numbers match.
	if lastRevision.UnlockConditions.UnlockHash() != contract.LastRevision.UnlockConditions.UnlockHash() {
		return errors.New("unlock conditions do not match")
	} else if lastRevision.NewRevisionNumber != contract.LastRevision.NewRevisionNumber {
		return &recentRevisionError{contract.LastRevision.NewRevisionNumber, lastRevision.NewRevisionNumber}
	}
	// NOTE: we can fake the blockheight here because it doesn't affect
	// verification; it just needs to be above the fork height and below the
	// contract expiration (which was checked earlier).
	return modules.VerifyFileContractRevisionTransactionSignatures(lastRevision, hostSignatures, contract.FileContract.WindowStart-1)
}

// negotiateRevision sends a revision and actions to the host for approval,
// completing one iteration of the revision loop.
func negotiateRevision(conn net.Conn, rev types.FileContractRevision, secretKey crypto.SecretKey) (types.Transaction, error) {
	// create transaction containing the revision
	signedTxn := types.Transaction{
		FileContractRevisions: []types.FileContractRevision{rev},
		TransactionSignatures: []types.TransactionSignature{{
			ParentID:       crypto.Hash(rev.ParentID),
			CoveredFields:  types.CoveredFields{FileContractRevisions: []uint64{0}},
			PublicKeyIndex: 0, // renter key is always first -- see formContract
		}},
	}
	// sign the transaction
	encodedSig, _ := crypto.SignHash(signedTxn.SigHash(0), secretKey) // no error possible
	signedTxn.TransactionSignatures[0].Signature = encodedSig[:]

	// send the revision
	if err := encoding.WriteObject(conn, rev); err != nil {
		return types.Transaction{}, errors.New("couldn't send revision: " + err.Error())
	}
	// read acceptance
	if err := modules.ReadNegotiationAcceptance(conn); err != nil {
		return types.Transaction{}, errors.New("host did not accept revision: " + err.Error())
	}

	// send the new transaction signature
	if err := encoding.WriteObject(conn, signedTxn.TransactionSignatures[0]); err != nil {
		return types.Transaction{}, errors.New("couldn't send transaction signature: " + err.Error())
	}
	// read the host's acceptance and transaction signature
	// NOTE: if the host sends ErrStopResponse, we should continue processing
	// the revision, but return the error anyway.
	responseErr := modules.ReadNegotiationAcceptance(conn)
	if responseErr != nil && responseErr != modules.ErrStopResponse {
		return types.Transaction{}, errors.New("host did not accept transaction signature: " + responseErr.Error())
	}
	var hostSig types.TransactionSignature
	if err := encoding.ReadObject(conn, &hostSig, 16e3); err != nil {
		return types.Transaction{}, errors.New("couldn't read host's signature: " + err.Error())
	}

	// add the signature to the transaction and verify it
	// NOTE: we can fake the blockheight here because it doesn't affect
	// verification; it just needs to be above the fork height and below the
	// contract expiration (which was checked earlier).
	verificationHeight := rev.NewWindowStart - 1
	signedTxn.TransactionSignatures = append(signedTxn.TransactionSignatures, hostSig)
	if err := signedTxn.StandaloneValid(verificationHeight); err != nil {
		return types.Transaction{}, err
	}

	// if the host sent ErrStopResponse, return it
	return signedTxn, responseErr
}

// newRevision creates a copy of current with its revision number incremented,
// and with cost transferred from the renter to the host.
func newRevision(current types.FileContractRevision, cost types.Currency) types.FileContractRevision {
	rev := current

	// need to manually copy slice memory
	rev.NewValidProofOutputs = make([]types.SiacoinOutput, 2)
	rev.NewMissedProofOutputs = make([]types.SiacoinOutput, 3)
	copy(rev.NewValidProofOutputs, current.NewValidProofOutputs)
	copy(rev.NewMissedProofOutputs, current.NewMissedProofOutputs)

	// move valid payout from renter to host
	rev.NewValidProofOutputs[0].Value = current.NewValidProofOutputs[0].Value.Sub(cost)
	rev.NewValidProofOutputs[1].Value = current.NewValidProofOutputs[1].Value.Add(cost)

	// move missed payout from renter to void
	rev.NewMissedProofOutputs[0].Value = current.NewMissedProofOutputs[0].Value.Sub(cost)
	rev.NewMissedProofOutputs[2].Value = current.NewMissedProofOutputs[2].Value.Add(cost)

	// increment revision number
	rev.NewRevisionNumber++

	return rev
}

// newDownloadRevision revises the current revision to cover the cost of
// downloading data.
func newDownloadRevision(current types.FileContractRevision, downloadCost types.Currency) types.FileContractRevision {
	return newRevision(current, downloadCost)
}

// newUploadRevision revises the current revision to cover the cost of
// uploading a sector.
func newUploadRevision(current types.FileContractRevision, merkleRoot crypto.Hash, price, collateral types.Currency) types.FileContractRevision {
	rev := newRevision(current, price)

	// move collateral from host to void
	rev.NewMissedProofOutputs[1].Value = rev.NewMissedProofOutputs[1].Value.Sub(collateral)
	rev.NewMissedProofOutputs[2].Value = rev.NewMissedProofOutputs[2].Value.Add(collateral)

	// set new filesize and Merkle root
	rev.NewFileSize += modules.SectorSize
	rev.NewFileMerkleRoot = merkleRoot
	return rev
}

// newDeleteRevision revises the current revision to cover the cost of
// deleting a sector.
func newDeleteRevision(current types.FileContractRevision, merkleRoot crypto.Hash) types.FileContractRevision {
	rev := newRevision(current, types.ZeroCurrency)
	rev.NewFileSize -= modules.SectorSize
	rev.NewFileMerkleRoot = merkleRoot
	return rev
}

// newModifyRevision revises the current revision to cover the cost of
// modifying a sector.
func newModifyRevision(current types.FileContractRevision, merkleRoot crypto.Hash, uploadCost types.Currency) types.FileContractRevision {
	rev := newRevision(current, uploadCost)
	rev.NewFileMerkleRoot = merkleRoot
	return rev
}
<|endoftext|>


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

package cassandra

import (
	"fmt"
	"sort"
	"testing"
	"time"

	"github.com/gocql/gocql"
	"github.com/square/metrics/api"
	"github.com/square/metrics/testing_support/assert"
)

func newDatabase(t *testing.T) *cassandraDatabase {
	cluster := gocql.NewCluster("localhost")
	cluster.Keyspace = "metrics_indexer_test"
	cluster.Consistency = gocql.One
	cluster.Timeout = time.Duration(10000 * time.Millisecond)
	session, err := cluster.CreateSession()
	if err != nil {
		t.Errorf("Cannot connect to Cassandra")
		return nil
	}
	tables := []string{"metric_names", "tag_index", "metric_name_set"}
	for _, table := range tables {
		if err := session.Query(fmt.Sprintf("TRUNCATE %s", table)).Exec(); err != nil {
			t.Errorf("Cannot truncate %s: %s", table, err.Error())
			return nil
		}
	}
	return &cassandraDatabase{
		session: session,
	}
}

func cleanDatabase(t *testing.T, db *cassandraDatabase) {
	db.session.Close()
}

func Test_MetricName_GetTagSet(t *testing.T) {
	a := assert.New(t)
	db := newDatabase(t)
	if db == nil {
		return
	}
	defer cleanDatabase(t, db)
	if db == nil {
		return
	}
	if _, err := db.GetTagSet("sample"); err == nil {
		t.Errorf("Cassandra should error on fetching nonexistent metric")
	}

	metricNamesTests := []struct {
		addTest      bool
		metricName   string
		tagString    string
		expectedTags map[string][]string // { metricName: [ tags ] }
	}{
		{true, "sample", "foo=bar1", map[string][]string{
			"sample": []string{"foo=bar1"},
		}},
		{true, "sample", "foo=bar2", map[string][]string{
			"sample": []string{"foo=bar1", "foo=bar2"},
		}},
		{true, "sample2", "foo=bar2", map[string][]string{
			"sample":  []string{"foo=bar1", "foo=bar2"},
			"sample2": []string{"foo=bar2"},
		}},
		{false, "sample2", "foo=bar2", map[string][]string{
			"sample": []string{"foo=bar1", "foo=bar2"},
		}},
		{false, "sample", "foo=bar1", map[string][]string{
			"sample": []string{"foo=bar2"},
		}},
	}

	for _, c := range metricNamesTests {
		if c.addTest {
			a.CheckError(db.AddMetricName(api.MetricKey(c.metricName), api.ParseTagSet(c.tagString)))
		} else {
			a.CheckError(db.RemoveMetricName(api.MetricKey(c.metricName), api.ParseTagSet(c.tagString)))
		}

		for k, v := range c.expectedTags {
			if tags, err := db.GetTagSet(api.MetricKey(k)); err != nil {
				t.Errorf("Error fetching tags")
			} else {
				stringTags := make([]string, len(tags))
				for i, tag := range tags {
					stringTags[i] = tag.Serialize()
				}

				a.EqInt(len(stringTags), len(v))
				sort.Sort(sort.StringSlice(stringTags))
				sort.Sort(sort.StringSlice(v))
				a.Eq(stringTags, v)
			}
		}
	}
}

func Test_GetAllMetrics(t *testing.T) {
	a := assert.New(t)
	db := newDatabase(t)
	if db == nil {
		return
	}
	defer cleanDatabase(t, db)
	a.CheckError(db.AddMetricName("metric.a", api.ParseTagSet("foo=a")))
	a.CheckError(db.AddMetricName("metric.a", api.ParseTagSet("foo=b")))
	a.CheckError(db.AddMetricNames([]api.TaggedMetric{
		{
			"metric.c",
			api.TagSet{
				"bar": "cat",
			},
		},
		{
			"metric.d",
			api.TagSet{
				"bar": "dog",
			},
		},
		{
			"metric.e",
			api.TagSet{
				"bar": "cat",
			},
		},
	}))
	keys, err := db.GetAllMetrics()
	a.CheckError(err)
	sort.Sort(api.MetricKeys(keys))
	a.Eq(keys, []api.MetricKey{"metric.a"})
	a.CheckError(db.AddMetricName("metric.b", api.ParseTagSet("foo=c")))
	a.CheckError(db.AddMetricName("metric.b", api.ParseTagSet("foo=c")))
	keys, err = db.GetAllMetrics()
	a.CheckError(err)
	sort.Sort(api.MetricKeys(keys))
	a.Eq(keys, []api.MetricKey{"metric.a", "metric.b", "metric.c", "metric.d", "metric.e"})

	lookupFooC, err := db.GetMetricKeys("foo", "c")
	a.CheckError(err)
	sort.Sort(api.MetricKeys(lookupFooC))
	a.Eq(lookupFooC, []api.MetricKey{"metric.b"})

	lookupBarCat, err := db.GetMetricKeys("bar", "cat")
	a.CheckError(err)
	sort.Sort(api.MetricKeys(lookupBarCat))
	a.Eq(lookupBarCat, []api.MetricKey{"metric.c", "metric.e"})
}

func Test_TagIndex(t *testing.T) {
	a := assert.New(t)
	db := newDatabase(t)
	if db == nil {
		return
	}
	defer cleanDatabase(t, db)

	if rows, err := db.GetMetricKeys("environment", "production"); err != nil {
		a.CheckError(err)
	} else {
		a.EqInt(len(rows), 0)
	}
	a.CheckError(db.AddToTagIndex("environment", "production", "a.b.c"))
	a.CheckError(db.AddToTagIndex("environment", "production", "d.e.f"))
	if rows, err := db.GetMetricKeys("environment", "production"); err != nil {
		a.CheckError(err)
	} else {
		a.EqInt(len(rows), 2)
	}

	a.CheckError(db.RemoveFromTagIndex("environment", "production", "a.b.c"))
	if rows, err := db.GetMetricKeys("environment", "production"); err != nil {
		a.CheckError(err)
	} else {
		a.EqInt(len(rows), 1)
		a.EqString(string(rows[0]), "d.e.f")
	}
}
<commit_msg>fix test<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cassandra

import (
	"fmt"
	"sort"
	"testing"
	"time"

	"github.com/gocql/gocql"
	"github.com/square/metrics/api"
	"github.com/square/metrics/testing_support/assert"
)

func newDatabase(t *testing.T) *cassandraDatabase {
	cluster := gocql.NewCluster("localhost")
	cluster.Keyspace = "metrics_indexer_test"
	cluster.Consistency = gocql.One
	cluster.Timeout = time.Duration(10000 * time.Millisecond)
	session, err := cluster.CreateSession()
	if err != nil {
		t.Errorf("Cannot connect to Cassandra")
		return nil
	}
	tables := []string{"metric_names", "tag_index", "metric_name_set"}
	for _, table := range tables {
		if err := session.Query(fmt.Sprintf("TRUNCATE %s", table)).Exec(); err != nil {
			t.Errorf("Cannot truncate %s: %s", table, err.Error())
			return nil
		}
	}
	return &cassandraDatabase{
		session: session,
	}
}

func cleanDatabase(t *testing.T, db *cassandraDatabase) {
	db.session.Close()
}

func Test_MetricName_GetTagSet(t *testing.T) {
	a := assert.New(t)
	db := newDatabase(t)
	if db == nil {
		return
	}
	defer cleanDatabase(t, db)
	if db == nil {
		return
	}
	if _, err := db.GetTagSet("sample"); err == nil {
		t.Errorf("Cassandra should error on fetching nonexistent metric")
	}

	metricNamesTests := []struct {
		addTest      bool
		metricName   string
		tagString    string
		expectedTags map[string][]string // { metricName: [ tags ] }
	}{
		{true, "sample", "foo=bar1", map[string][]string{
			"sample": []string{"foo=bar1"},
		}},
		{true, "sample", "foo=bar2", map[string][]string{
			"sample": []string{"foo=bar1", "foo=bar2"},
		}},
		{true, "sample2", "foo=bar2", map[string][]string{
			"sample":  []string{"foo=bar1", "foo=bar2"},
			"sample2": []string{"foo=bar2"},
		}},
		{false, "sample2", "foo=bar2", map[string][]string{
			"sample": []string{"foo=bar1", "foo=bar2"},
		}},
		{false, "sample", "foo=bar1", map[string][]string{
			"sample": []string{"foo=bar2"},
		}},
	}

	for _, c := range metricNamesTests {
		if c.addTest {
			a.CheckError(db.AddMetricName(api.MetricKey(c.metricName), api.ParseTagSet(c.tagString)))
		} else {
			a.CheckError(db.RemoveMetricName(api.MetricKey(c.metricName), api.ParseTagSet(c.tagString)))
		}

		for k, v := range c.expectedTags {
			if tags, err := db.GetTagSet(api.MetricKey(k)); err != nil {
				t.Errorf("Error fetching tags")
			} else {
				stringTags := make([]string, len(tags))
				for i, tag := range tags {
					stringTags[i] = tag.Serialize()
				}

				a.EqInt(len(stringTags), len(v))
				sort.Sort(sort.StringSlice(stringTags))
				sort.Sort(sort.StringSlice(v))
				a.Eq(stringTags, v)
			}
		}
	}
}

func Test_GetAllMetrics(t *testing.T) {
	a := assert.New(t)
	db := newDatabase(t)
	if db == nil {
		return
	}
	defer cleanDatabase(t, db)
	a.CheckError(db.AddMetricName("metric.a", api.ParseTagSet("foo=a")))
	a.CheckError(db.AddMetricName("metric.a", api.ParseTagSet("foo=b")))
	a.CheckError(db.AddMetricNames([]api.TaggedMetric{
		{
			"metric.c",
			api.TagSet{
				"bar": "cat",
			},
		},
		{
			"metric.d",
			api.TagSet{
				"bar": "dog",
			},
		},
		{
			"metric.e",
			api.TagSet{
				"bar": "cat",
			},
		},
	}))
	keys, err := db.GetAllMetrics()
	a.CheckError(err)
	sort.Sort(api.MetricKeys(keys))
	a.Eq(keys, []api.MetricKey{"metric.a", "metric.c", "metric.d", "metric.e"})
	a.CheckError(db.AddMetricName("metric.b", api.ParseTagSet("foo=c")))
	a.CheckError(db.AddMetricName("metric.b", api.ParseTagSet("foo=c")))
	keys, err = db.GetAllMetrics()
	a.CheckError(err)
	sort.Sort(api.MetricKeys(keys))
	a.Eq(keys, []api.MetricKey{"metric.a", "metric.b", "metric.c", "metric.d", "metric.e"})

	lookupFooC, err := db.GetMetricKeys("foo", "c")
	a.CheckError(err)
	sort.Sort(api.MetricKeys(lookupFooC))
	a.Eq(lookupFooC, []api.MetricKey{"metric.b"})

	lookupBarCat, err := db.GetMetricKeys("bar", "cat")
	a.CheckError(err)
	sort.Sort(api.MetricKeys(lookupBarCat))
	a.Eq(lookupBarCat, []api.MetricKey{"metric.c", "metric.e"})
}

func Test_TagIndex(t *testing.T) {
	a := assert.New(t)
	db := newDatabase(t)
	if db == nil {
		return
	}
	defer cleanDatabase(t, db)

	if rows, err := db.GetMetricKeys("environment", "production"); err != nil {
		a.CheckError(err)
	} else {
		a.EqInt(len(rows), 0)
	}
	a.CheckError(db.AddToTagIndex("environment", "production", "a.b.c"))
	a.CheckError(db.AddToTagIndex("environment", "production", "d.e.f"))
	if rows, err := db.GetMetricKeys("environment", "production"); err != nil {
		a.CheckError(err)
	} else {
		a.EqInt(len(rows), 2)
	}

	a.CheckError(db.RemoveFromTagIndex("environment", "production", "a.b.c"))
	if rows, err := db.GetMetricKeys("environment", "production"); err != nil {
		a.CheckError(err)
	} else {
		a.EqInt(len(rows), 1)
		a.EqString(string(rows[0]), "d.e.f")
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"fmt"

	"github.com/docker/distribution"
	"github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/manifest/schema1"
	"github.com/docker/distribution/reference"
	"github.com/docker/libtrust"
)

type manifestStore struct {
	repository                 *repository
	revisionStore              *revisionStore
	tagStore                   *tagStore
	ctx                        context.Context
	skipDependencyVerification bool
}

var _ distribution.ManifestService = &manifestStore{}

func (ms *manifestStore) Exists(dgst digest.Digest) (bool, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Exists")

	_, err := ms.revisionStore.blobStore.Stat(ms.ctx, dgst)
	if err != nil {
		if err == distribution.ErrBlobUnknown {
			return false, nil
		}

		return false, err
	}

	return true, nil
}

func (ms *manifestStore) Get(dgst digest.Digest) (*schema1.SignedManifest, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Get")
	return ms.revisionStore.get(ms.ctx, dgst)
}

// SkipLayerVerification allows a manifest to be Put before it's
// layers are on the filesystem
func SkipLayerVerification(ms distribution.ManifestService) error {
	if ms, ok := ms.(*manifestStore); ok {
		ms.skipDependencyVerification = true
		return nil
	}
	return fmt.Errorf("skip layer verification only valid for manifestStore")
}

func (ms *manifestStore) Put(manifest *schema1.SignedManifest) error {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Put")

	if err := ms.verifyManifest(ms.ctx, manifest); err != nil {
		return err
	}

	// Store the revision of the manifest
	revision, err := ms.revisionStore.put(ms.ctx, manifest)
	if err != nil {
		return err
	}

	// Now, tag the manifest
	return ms.tagStore.tag(manifest.Tag, revision.Digest)
}

// Delete removes the revision of the specified manfiest.
func (ms *manifestStore) Delete(dgst digest.Digest) error {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Delete")
	return ms.revisionStore.delete(ms.ctx, dgst)
}

func (ms *manifestStore) Tags() ([]string, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Tags")
	return ms.tagStore.tags()
}

func (ms *manifestStore) ExistsByTag(tag string) (bool, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).ExistsByTag")
	return ms.tagStore.exists(tag)
}

func (ms *manifestStore) GetByTag(tag string, options ...distribution.ManifestServiceOption) (*schema1.SignedManifest, error) {
	for _, option := range options {
		err := option(ms)
		if err != nil {
			return nil, err
		}
	}

	context.GetLogger(ms.ctx).Debug("(*manifestStore).GetByTag")
	dgst, err := ms.tagStore.resolve(tag)
	if err != nil {
		return nil, err
	}

	return ms.revisionStore.get(ms.ctx, dgst)
}

// verifyManifest ensures that the manifest content is valid from the
// perspective of the registry. It ensures that the signature is valid for the
// enclosed payload. As a policy, the registry only tries to store valid
// content, leaving trust policies of that content up to consumers.
func (ms *manifestStore) verifyManifest(ctx context.Context, mnfst *schema1.SignedManifest) error {
	var errs distribution.ErrManifestVerification

	if len(mnfst.Name) > reference.NameTotalLengthMax {
		errs = append(errs, fmt.Errorf("manifest name must not be more than %v characters", reference.NameTotalLengthMax))
	}

	if len(mnfst.History) != len(mnfst.FSLayers) {
		errs = append(errs, fmt.Errorf("mismatched history and fslayer cardinality %d != %d",
			len(mnfst.History), len(mnfst.FSLayers)))
	}

	if _, err := schema1.Verify(mnfst); err != nil {
		switch err {
		case libtrust.ErrMissingSignatureKey, libtrust.ErrInvalidJSONContent, libtrust.ErrMissingSignatureKey:
			errs = append(errs, distribution.ErrManifestUnverified{})
		default:
			if err.Error() == "invalid signature" { // TODO(stevvooe): This should be exported by libtrust
				errs = append(errs, distribution.ErrManifestUnverified{})
			} else {
				errs = append(errs, err)
			}
		}
	}

	if !ms.skipDependencyVerification {
		for _, fsLayer := range mnfst.FSLayers {
			_, err := ms.repository.Blobs(ctx).Stat(ctx, fsLayer.BlobSum)
			if err != nil {
				if err != distribution.ErrBlobUnknown {
					errs = append(errs, err)
				}

				// On error here, we always append unknown blob errors.
				errs = append(errs, distribution.ErrManifestBlobUnknown{Digest: fsLayer.BlobSum})
			}
		}
	}
	if len(errs) != 0 {
		return errs
	}

	return nil
}
<commit_msg>Verify manifest name format<commit_after>package storage

import (
	"fmt"

	"github.com/docker/distribution"
	"github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/manifest/schema1"
	"github.com/docker/distribution/reference"
	"github.com/docker/libtrust"
)

type manifestStore struct {
	repository                 *repository
	revisionStore              *revisionStore
	tagStore                   *tagStore
	ctx                        context.Context
	skipDependencyVerification bool
}

var _ distribution.ManifestService = &manifestStore{}

func (ms *manifestStore) Exists(dgst digest.Digest) (bool, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Exists")

	_, err := ms.revisionStore.blobStore.Stat(ms.ctx, dgst)
	if err != nil {
		if err == distribution.ErrBlobUnknown {
			return false, nil
		}

		return false, err
	}

	return true, nil
}

func (ms *manifestStore) Get(dgst digest.Digest) (*schema1.SignedManifest, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Get")
	return ms.revisionStore.get(ms.ctx, dgst)
}

// SkipLayerVerification allows a manifest to be Put before it's
// layers are on the filesystem
func SkipLayerVerification(ms distribution.ManifestService) error {
	if ms, ok := ms.(*manifestStore); ok {
		ms.skipDependencyVerification = true
		return nil
	}
	return fmt.Errorf("skip layer verification only valid for manifestStore")
}

func (ms *manifestStore) Put(manifest *schema1.SignedManifest) error {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Put")

	if err := ms.verifyManifest(ms.ctx, manifest); err != nil {
		return err
	}

	// Store the revision of the manifest
	revision, err := ms.revisionStore.put(ms.ctx, manifest)
	if err != nil {
		return err
	}

	// Now, tag the manifest
	return ms.tagStore.tag(manifest.Tag, revision.Digest)
}

// Delete removes the revision of the specified manfiest.
func (ms *manifestStore) Delete(dgst digest.Digest) error {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Delete")
	return ms.revisionStore.delete(ms.ctx, dgst)
}

func (ms *manifestStore) Tags() ([]string, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).Tags")
	return ms.tagStore.tags()
}

func (ms *manifestStore) ExistsByTag(tag string) (bool, error) {
	context.GetLogger(ms.ctx).Debug("(*manifestStore).ExistsByTag")
	return ms.tagStore.exists(tag)
}

func (ms *manifestStore) GetByTag(tag string, options ...distribution.ManifestServiceOption) (*schema1.SignedManifest, error) {
	for _, option := range options {
		err := option(ms)
		if err != nil {
			return nil, err
		}
	}

	context.GetLogger(ms.ctx).Debug("(*manifestStore).GetByTag")
	dgst, err := ms.tagStore.resolve(tag)
	if err != nil {
		return nil, err
	}

	return ms.revisionStore.get(ms.ctx, dgst)
}

// verifyManifest ensures that the manifest content is valid from the
// perspective of the registry. It ensures that the signature is valid for the
// enclosed payload. As a policy, the registry only tries to store valid
// content, leaving trust policies of that content up to consumers.
func (ms *manifestStore) verifyManifest(ctx context.Context, mnfst *schema1.SignedManifest) error {
	var errs distribution.ErrManifestVerification

	if len(mnfst.Name) > reference.NameTotalLengthMax {
		errs = append(errs, fmt.Errorf("manifest name must not be more than %v characters", reference.NameTotalLengthMax))
	}

	if !reference.NameRegexp.MatchString(mnfst.Name) {
		errs = append(errs, fmt.Errorf("invalid manifest name format"))
	}

	if len(mnfst.History) != len(mnfst.FSLayers) {
		errs = append(errs, fmt.Errorf("mismatched history and fslayer cardinality %d != %d",
			len(mnfst.History), len(mnfst.FSLayers)))
	}

	if _, err := schema1.Verify(mnfst); err != nil {
		switch err {
		case libtrust.ErrMissingSignatureKey, libtrust.ErrInvalidJSONContent, libtrust.ErrMissingSignatureKey:
			errs = append(errs, distribution.ErrManifestUnverified{})
		default:
			if err.Error() == "invalid signature" { // TODO(stevvooe): This should be exported by libtrust
				errs = append(errs, distribution.ErrManifestUnverified{})
			} else {
				errs = append(errs, err)
			}
		}
	}

	if !ms.skipDependencyVerification {
		for _, fsLayer := range mnfst.FSLayers {
			_, err := ms.repository.Blobs(ctx).Stat(ctx, fsLayer.BlobSum)
			if err != nil {
				if err != distribution.ErrBlobUnknown {
					errs = append(errs, err)
				}

				// On error here, we always append unknown blob errors.
				errs = append(errs, distribution.ErrManifestBlobUnknown{Digest: fsLayer.BlobSum})
			}
		}
	}
	if len(errs) != 0 {
		return errs
	}

	return nil
}
<|endoftext|>


<commit_before>package service

import (
	"time"

	commonModel "github.com/Cepave/open-falcon-backend/common/model"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/modules/hbs/cache"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/model"
	"github.com/icrowley/fake"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var dummyTime = time.Now().Unix()

type fakeHeartbeat struct {
	rowsAffectedCnt int
	alwaysSuccess   bool
}

func (a *fakeHeartbeat) calling(agents []*model.AgentHeartbeat) (int64, int64) {
	count := len(agents)
	a.rowsAffectedCnt += count

	if a.alwaysSuccess {
		return int64(count), 0
	} else {
		return 0, int64(count)
	}
}

func generateRandomHeartbeat() *commonModel.AgentReportRequest {
	return &commonModel.AgentReportRequest{
		Hostname:     fake.DomainName(),
		IP:           fake.IPv4(),
		AgentVersion: fake.Digits(),
	}
}

func eventuallyWithTimeout(valueGetter interface{}, timeout time.Duration) GomegaAsyncAssertion {
	return Eventually(
		valueGetter, timeout, timeout/8,
	)
}

var _ = Describe("Test the behavior of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
		heartbeatImpl         *fakeHeartbeat
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{Num: 16, Dur: 100 * time.Millisecond},
		)
		heartbeatImpl = &fakeHeartbeat{alwaysSuccess: true}
		agentHeartbeatService.heartbeatCall = heartbeatImpl.calling
	})

	It("should work normally during its life cycle", func() {
		batchSize := agentHeartbeatService.qConfig.Num
		dataNum := batchSize * 2

		agentHeartbeatService.Start()
		for i := 0; i < dataNum; i++ {
			agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
		}
		agentHeartbeatService.Stop()

		Expect(agentHeartbeatService.CumulativeAgentsPut()).To(Equal(int64(dataNum)))
		eventuallyWithTimeout(func() int {
			return heartbeatImpl.rowsAffectedCnt
		}, time.Second).Should(Equal(dataNum))
		eventuallyWithTimeout(func() int {
			return agentHeartbeatService.CurrentSize()
		}, time.Second).Should(BeZero())
	})
})

var _ = Describe("Test Put() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
		heartbeatImpl         *fakeHeartbeat
		now                   int64
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{Num: 16},
		)
		heartbeatImpl = &fakeHeartbeat{alwaysSuccess: true}
		agentHeartbeatService.heartbeatCall = heartbeatImpl.calling
		now = time.Now().Unix()
	})

	Context("when service is not running", func() {
		It("should not add data", func() {
			data := generateRandomHeartbeat()
			agentHeartbeatService.Put(data, now)
			_, ok := cache.Agents.Get(data.Hostname)

			Expect(ok).To(BeFalse())
			Expect(agentHeartbeatService.CumulativeAgentsPut()).To(BeZero())
			Expect(agentHeartbeatService.CurrentSize()).To(BeZero())
		})
	})

	Context("when service is running", func() {
		It("should add data", func() {
			data := generateRandomHeartbeat()
			agentHeartbeatService.running = true
			agentHeartbeatService.Put(data, now)
			val, ok := cache.Agents.Get(data.Hostname)

			Expect(ok).To(BeTrue())
			Expect(val.ReportRequest.Hostname).To(Equal(data.Hostname))
			Expect(val.LastUpdate).To(Equal(now))
			Expect(agentHeartbeatService.CumulativeAgentsPut()).To(Equal(int64(1)))
			Expect(agentHeartbeatService.CurrentSize()).To(Equal(1))
		})
	})
})

var _ = Describe("Test Start() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{},
		)
	})

	AfterEach(func() {
		agentHeartbeatService.Stop()
	})

	Context("when service is stopped", func() {
		It("Start() should change the running status", func() {
			agentHeartbeatService.Start()
			Expect(agentHeartbeatService.running).To(BeTrue())
		})
	})

	Context("when service is started", func() {
		It("Start() should not change the running status", func() {
			agentHeartbeatService.running = true
			agentHeartbeatService.Start()
			Expect(agentHeartbeatService.running).To(BeTrue())
		})
	})
})

var _ = Describe("Test Stop() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{},
		)
	})

	Context("when service is stopped", func() {
		It("Stop() should not change the running status", func() {
			agentHeartbeatService.Stop()
			Expect(agentHeartbeatService.running).To(BeFalse())
		})
	})

	Context("when service is started", func() {
		It("Stop() should change the running status", func() {
			agentHeartbeatService.running = true
			agentHeartbeatService.Stop()
			Expect(agentHeartbeatService.running).To(BeFalse())
		})
	})
})

var _ = Describe("Test consumeHeartbeatQueue() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
		heartbeatImpl         *fakeHeartbeat
	)

	BeforeEach(func() {
		heartbeatImpl = &fakeHeartbeat{alwaysSuccess: true}
	})

	JustBeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{Num: 16},
		)
		agentHeartbeatService.heartbeatCall = heartbeatImpl.calling
		agentHeartbeatService.running = true
	})

	Context("when success", func() {
		It("rowsAffectedCnt should be incremented normally", func() {
			agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			agentHeartbeatService.consumeHeartbeatQueue(false)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(1))
			Expect(agentHeartbeatService.CumulativeRowsAffected()).To(Equal(int64(1)))
			Expect(agentHeartbeatService.CumulativeAgentsDropped()).To(BeZero())
		})
	})

	Context("when failure", func() {
		BeforeEach(func() {
			heartbeatImpl = &fakeHeartbeat{alwaysSuccess: false}
		})

		It("agentsDroppedCnt should be incremented normally", func() {
			agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			agentHeartbeatService.consumeHeartbeatQueue(false)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(1))
			Expect(agentHeartbeatService.CumulativeRowsAffected()).To(BeZero())
			Expect(agentHeartbeatService.CumulativeAgentsDropped()).To(Equal(int64(1)))
		})
	})

	Context("when non-flushing mode", func() {
		It("should consume an amount of data = batch size", func() {
			batchSize := agentHeartbeatService.qConfig.Num
			dataNum := batchSize*2 - 1
			for i := 0; i < dataNum; i++ {
				agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			}
			agentHeartbeatService.consumeHeartbeatQueue(false)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(batchSize))
			Expect(agentHeartbeatService.CurrentSize()).To(Equal(batchSize - 1))
		})
	})

	Context("when flushing mode", func() {
		It("should flush data to 0", func() {
			dataNum := agentHeartbeatService.qConfig.Num*2 - 1
			for i := 0; i < dataNum; i++ {
				agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			}
			agentHeartbeatService.consumeHeartbeatQueue(true)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(dataNum))
			Expect(agentHeartbeatService.CurrentSize()).To(BeZero())
		})
	})
})
<commit_msg>[OWL-1695] Clean agents in cache package after testing.<commit_after>package service

import (
	"time"

	commonModel "github.com/Cepave/open-falcon-backend/common/model"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/modules/hbs/cache"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/model"
	"github.com/icrowley/fake"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var dummyTime = time.Now().Unix()

type fakeHeartbeat struct {
	rowsAffectedCnt int
	alwaysSuccess   bool
}

func (a *fakeHeartbeat) calling(agents []*model.AgentHeartbeat) (int64, int64) {
	count := len(agents)
	a.rowsAffectedCnt += count

	if a.alwaysSuccess {
		return int64(count), 0
	} else {
		return 0, int64(count)
	}
}

func generateRandomHeartbeat() *commonModel.AgentReportRequest {
	return &commonModel.AgentReportRequest{
		Hostname:     fake.DomainName(),
		IP:           fake.IPv4(),
		AgentVersion: fake.Digits(),
	}
}

func eventuallyWithTimeout(valueGetter interface{}, timeout time.Duration) GomegaAsyncAssertion {
	return Eventually(
		valueGetter, timeout, timeout/8,
	)
}

var _ = Describe("Test the behavior of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
		heartbeatImpl         *fakeHeartbeat
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{Num: 16, Dur: 100 * time.Millisecond},
		)
		heartbeatImpl = &fakeHeartbeat{alwaysSuccess: true}
		agentHeartbeatService.heartbeatCall = heartbeatImpl.calling
	})

	It("should work normally during its life cycle", func() {
		batchSize := agentHeartbeatService.qConfig.Num
		dataNum := batchSize * 2

		agentHeartbeatService.Start()
		for i := 0; i < dataNum; i++ {
			agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
		}
		agentHeartbeatService.Stop()

		Expect(agentHeartbeatService.CumulativeAgentsPut()).To(Equal(int64(dataNum)))
		eventuallyWithTimeout(func() int {
			return heartbeatImpl.rowsAffectedCnt
		}, time.Second).Should(Equal(dataNum))
		eventuallyWithTimeout(func() int {
			return agentHeartbeatService.CurrentSize()
		}, time.Second).Should(BeZero())
	})
})

var _ = Describe("Test Put() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
		heartbeatImpl         *fakeHeartbeat
		now                   int64
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{Num: 16},
		)
		heartbeatImpl = &fakeHeartbeat{alwaysSuccess: true}
		agentHeartbeatService.heartbeatCall = heartbeatImpl.calling
		now = time.Now().Unix()
	})

	AfterEach(func() {
		cache.Agents = cache.NewSafeAgents()
	})

	Context("when service is not running", func() {
		It("should not add data", func() {
			data := generateRandomHeartbeat()
			agentHeartbeatService.Put(data, now)
			_, ok := cache.Agents.Get(data.Hostname)

			Expect(ok).To(BeFalse())
			Expect(agentHeartbeatService.CumulativeAgentsPut()).To(BeZero())
			Expect(agentHeartbeatService.CurrentSize()).To(BeZero())
		})
	})

	Context("when service is running", func() {
		It("should add data", func() {
			data := generateRandomHeartbeat()
			agentHeartbeatService.running = true
			agentHeartbeatService.Put(data, now)
			val, ok := cache.Agents.Get(data.Hostname)

			Expect(ok).To(BeTrue())
			Expect(val.ReportRequest.Hostname).To(Equal(data.Hostname))
			Expect(val.LastUpdate).To(Equal(now))
			Expect(agentHeartbeatService.CumulativeAgentsPut()).To(Equal(int64(1)))
			Expect(agentHeartbeatService.CurrentSize()).To(Equal(1))
		})
	})
})

var _ = Describe("Test Start() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{},
		)
	})

	AfterEach(func() {
		agentHeartbeatService.Stop()
	})

	Context("when service is stopped", func() {
		It("Start() should change the running status", func() {
			agentHeartbeatService.Start()
			Expect(agentHeartbeatService.running).To(BeTrue())
		})
	})

	Context("when service is started", func() {
		It("Start() should not change the running status", func() {
			agentHeartbeatService.running = true
			agentHeartbeatService.Start()
			Expect(agentHeartbeatService.running).To(BeTrue())
		})
	})
})

var _ = Describe("Test Stop() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
	)

	BeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{},
		)
	})

	Context("when service is stopped", func() {
		It("Stop() should not change the running status", func() {
			agentHeartbeatService.Stop()
			Expect(agentHeartbeatService.running).To(BeFalse())
		})
	})

	Context("when service is started", func() {
		It("Stop() should change the running status", func() {
			agentHeartbeatService.running = true
			agentHeartbeatService.Stop()
			Expect(agentHeartbeatService.running).To(BeFalse())
		})
	})
})

var _ = Describe("Test consumeHeartbeatQueue() of AgentHeartbeat service", func() {
	var (
		agentHeartbeatService *AgentHeartbeatService
		heartbeatImpl         *fakeHeartbeat
	)

	BeforeEach(func() {
		heartbeatImpl = &fakeHeartbeat{alwaysSuccess: true}
	})

	JustBeforeEach(func() {
		agentHeartbeatService = NewAgentHeartbeatService(
			&commonQueue.Config{Num: 16},
		)
		agentHeartbeatService.heartbeatCall = heartbeatImpl.calling
		agentHeartbeatService.running = true
	})

	Context("when success", func() {
		It("rowsAffectedCnt should be incremented normally", func() {
			agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			agentHeartbeatService.consumeHeartbeatQueue(false)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(1))
			Expect(agentHeartbeatService.CumulativeRowsAffected()).To(Equal(int64(1)))
			Expect(agentHeartbeatService.CumulativeAgentsDropped()).To(BeZero())
		})
	})

	Context("when failure", func() {
		BeforeEach(func() {
			heartbeatImpl = &fakeHeartbeat{alwaysSuccess: false}
		})

		It("agentsDroppedCnt should be incremented normally", func() {
			agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			agentHeartbeatService.consumeHeartbeatQueue(false)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(1))
			Expect(agentHeartbeatService.CumulativeRowsAffected()).To(BeZero())
			Expect(agentHeartbeatService.CumulativeAgentsDropped()).To(Equal(int64(1)))
		})
	})

	Context("when non-flushing mode", func() {
		It("should consume an amount of data = batch size", func() {
			batchSize := agentHeartbeatService.qConfig.Num
			dataNum := batchSize*2 - 1
			for i := 0; i < dataNum; i++ {
				agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			}
			agentHeartbeatService.consumeHeartbeatQueue(false)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(batchSize))
			Expect(agentHeartbeatService.CurrentSize()).To(Equal(batchSize - 1))
		})
	})

	Context("when flushing mode", func() {
		It("should flush data to 0", func() {
			dataNum := agentHeartbeatService.qConfig.Num*2 - 1
			for i := 0; i < dataNum; i++ {
				agentHeartbeatService.Put(generateRandomHeartbeat(), dummyTime)
			}
			agentHeartbeatService.consumeHeartbeatQueue(true)

			Expect(heartbeatImpl.rowsAffectedCnt).To(Equal(dataNum))
			Expect(agentHeartbeatService.CurrentSize()).To(BeZero())
		})
	})
})
<|endoftext|>


<commit_before>package disambig

import (
	. "chukuparser/nlp/types"
)

func Full(s Spellout) string {
	return s.AsString()
}
<commit_msg>Added more parameter function types<commit_after>package disambig

import (
	. "chukuparser/nlp/types"
	"fmt"
)

const (
	SEPARATOR = ";"
)

var Main_POS map[string]bool

func init() {
	Main_POS_Types := []string{"NN", "VB", "RR", "VB"}
	Main_POS = make(map[string]bool, len(Main_POS_Types))
	for _, pos := range Main_POS_Types {
		Main_POS[pos] = true
	}
}

func Full(s Spellout) string {
	return s.AsString()
}

type MProject func(m Morpheme) string

func projectMorphemes(s Spellout, f MProject) string {
	strs := make([]string, len(s))
	for i, morph := range s {
		strs[i] = MProject(morph)
	}
	return strings.Join(strs, SEPARATOR)
}

func Segments(s Spellout) string {
	return projectMorphemes(s, func(m Morpheme) {
		return m.Form
	})
}

func POS_Props(s Spellout) string {
	return projectMorphemes(s, func(m Morpheme) {
		return fmt.Sprintf("%s_%s", m.POS, m.FeatureStr)
	})
}

func Funcs_Main_POS_Props(s Spellout) string {
	strs := make([]string, len(s))
	var exists bool
	for i, morph := range s {
		_, exists = Main_POS[morph.POS]
		if exists {
			strs[i] = fmt.Sprintf("%s_%s", m.POS, m.FeatureStr)
		} else {
			strs[i] = morph.Form
		}
	}
	return strings.Join(strs, SEPARATOR)
}

func POS(s Spellout) string {
	return projectMorphemes(s, func(m Morpheme) {
		return m.POS
	})
}
<|endoftext|>


<commit_before>package datasource

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

var expBackoffTests = []struct {
	count int
	body  string
}{
	{0, "number of attempts: 0"},
	{1, "number of attempts: 1"},
	{2, "number of attempts: 2"},
}

// Test exponential backoff and that it continues retrying if a 5xx response is
// received
func TestFetchURLExpBackOff(t *testing.T) {
	for i, tt := range expBackoffTests {
		mux := http.NewServeMux()
		count := 0
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			if count == tt.count {
				io.WriteString(w, fmt.Sprintf("number of attempts: %d", count))
				return
			}
			count++
			http.Error(w, "", 500)
		})
		ts := httptest.NewServer(mux)
		defer ts.Close()

		data, err := fetchURL(ts.URL)
		if err != nil {
			t.Errorf("Test case %d produced error: %v", i, err)
		}

		if count != tt.count {
			t.Errorf("Test case %d failed: %d != %d", i, count, tt.count)
		}

		if string(data) != tt.body {
			t.Errorf("Test case %d failed: %s != %s", i, tt.body, data)
		}
	}
}

// Test that it stops retrying if a 4xx response comes back
func TestFetchURL4xx(t *testing.T) {
	retries := 0
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		retries++
		http.Error(w, "", 404)
	}))
	defer ts.Close()

	_, err := fetchURL(ts.URL)
	if err == nil {
		t.Errorf("Incorrect result\ngot:  %s\nwant: %s", err.Error(), "user-data not found. HTTP status code: 404")
	}

	if retries > 1 {
		t.Errorf("Number of retries:\n%d\nExpected number of retries:\n%s", retries, 1)
	}
}

// Test that it fetches and returns user-data just fine
func TestFetchURL2xx(t *testing.T) {
	var cloudcfg = `
#cloud-config
coreos:
	oem:
	    id: test
	    name: CoreOS.box for Test
	    version-id: %VERSION_ID%+%BUILD_ID%
	    home-url: https://github.com/coreos/coreos-cloudinit
	    bug-report-url: https://github.com/coreos/coreos-cloudinit
	update:
		reboot-strategy: best-effort
`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, cloudcfg)
	}))
	defer ts.Close()

	data, err := fetchURL(ts.URL)
	if err != nil {
		t.Errorf("Incorrect result\ngot:  %v\nwant: %v", err, nil)
	}

	if string(data) != cloudcfg {
		t.Errorf("Incorrect result\ngot:  %s\nwant: %s", string(data), cloudcfg)
	}
}

// Test attempt to fetching using malformed URL
func TestFetchURLMalformed(t *testing.T) {
	var tests = []struct {
		url  string
		want string
	}{
		{"boo", "user-data URL boo does not have a valid HTTP scheme. Skipping."},
		{"mailto://boo", "user-data URL mailto://boo does not have a valid HTTP scheme. Skipping."},
		{"ftp://boo", "user-data URL ftp://boo does not have a valid HTTP scheme. Skipping."},
		{"", "user-data URL is empty. Skipping."},
	}

	for _, test := range tests {
		_, err := fetchURL(test.url)
		if err != nil && err.Error() != test.want {
			t.Errorf("Incorrect result\ngot:  %v\nwant: %v", err, test.want)
		}
	}
}
<commit_msg>test(datastore/fetch): Makes sure err is not nil<commit_after>package datasource

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

var expBackoffTests = []struct {
	count int
	body  string
}{
	{0, "number of attempts: 0"},
	{1, "number of attempts: 1"},
	{2, "number of attempts: 2"},
}

// Test exponential backoff and that it continues retrying if a 5xx response is
// received
func TestFetchURLExpBackOff(t *testing.T) {
	for i, tt := range expBackoffTests {
		mux := http.NewServeMux()
		count := 0
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			if count == tt.count {
				io.WriteString(w, fmt.Sprintf("number of attempts: %d", count))
				return
			}
			count++
			http.Error(w, "", 500)
		})
		ts := httptest.NewServer(mux)
		defer ts.Close()

		data, err := fetchURL(ts.URL)
		if err != nil {
			t.Errorf("Test case %d produced error: %v", i, err)
		}

		if count != tt.count {
			t.Errorf("Test case %d failed: %d != %d", i, count, tt.count)
		}

		if string(data) != tt.body {
			t.Errorf("Test case %d failed: %s != %s", i, tt.body, data)
		}
	}
}

// Test that it stops retrying if a 4xx response comes back
func TestFetchURL4xx(t *testing.T) {
	retries := 0
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		retries++
		http.Error(w, "", 404)
	}))
	defer ts.Close()

	_, err := fetchURL(ts.URL)
	if err == nil {
		t.Errorf("Incorrect result\ngot:  %s\nwant: %s", err.Error(), "user-data not found. HTTP status code: 404")
	}

	if retries > 1 {
		t.Errorf("Number of retries:\n%d\nExpected number of retries:\n%s", retries, 1)
	}
}

// Test that it fetches and returns user-data just fine
func TestFetchURL2xx(t *testing.T) {
	var cloudcfg = `
#cloud-config
coreos:
	oem:
	    id: test
	    name: CoreOS.box for Test
	    version-id: %VERSION_ID%+%BUILD_ID%
	    home-url: https://github.com/coreos/coreos-cloudinit
	    bug-report-url: https://github.com/coreos/coreos-cloudinit
	update:
		reboot-strategy: best-effort
`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, cloudcfg)
	}))
	defer ts.Close()

	data, err := fetchURL(ts.URL)
	if err != nil {
		t.Errorf("Incorrect result\ngot:  %v\nwant: %v", err, nil)
	}

	if string(data) != cloudcfg {
		t.Errorf("Incorrect result\ngot:  %s\nwant: %s", string(data), cloudcfg)
	}
}

// Test attempt to fetching using malformed URL
func TestFetchURLMalformed(t *testing.T) {
	var tests = []struct {
		url  string
		want string
	}{
		{"boo", "user-data URL boo does not have a valid HTTP scheme. Skipping."},
		{"mailto://boo", "user-data URL mailto://boo does not have a valid HTTP scheme. Skipping."},
		{"ftp://boo", "user-data URL ftp://boo does not have a valid HTTP scheme. Skipping."},
		{"", "user-data URL is empty. Skipping."},
	}

	for _, test := range tests {
		_, err := fetchURL(test.url)
		if err == nil || err.Error() != test.want {
			t.Errorf("Incorrect result\ngot:  %v\nwant: %v", err, test.want)
		}
	}
}
<|endoftext|>


<commit_before>package zoneFileParser

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

	"github.com/netsec-ethz/rains/rainslib"

	log "github.com/inconshreveable/log15"
	"golang.org/x/crypto/ed25519"
)

//encodeZone return z in zonefile format. If addZoneAndContext is true, the context and subject zone
//are present also for all contained sections.
func encodeZone(z *rainslib.ZoneSection, addZoneAndContext bool) string {
	zone := fmt.Sprintf("%s %s %s [\n", TypeZone, z.SubjectZone, z.Context)
	for _, section := range z.Content {
		switch section := section.(type) {
		case *rainslib.AssertionSection:
			zone += encodeAssertion(section, z.Context, z.SubjectZone, indent4, addZoneAndContext)
		case *rainslib.ShardSection:
			zone += encodeShard(section, z.Context, z.SubjectZone, indent4, addZoneAndContext)
		default:
			log.Warn("Unsupported message section type", "msgSection", section)
		}
	}
	if z.Signatures != nil {
		var sigs []string
		for _, sig := range z.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			return fmt.Sprintf("%s] ( %s )\n", zone, sigs[0])
		}
		return fmt.Sprintf("%s] ( \n%s%s\n  )\n", zone, indent4, strings.Join(sigs, "\n"+indent4))
	}
	return fmt.Sprintf("%s]\n", zone)
}

//encodeShard returns s in zonefile format. If addZoneAndContext is true, the context and subject
//zone are present for the shard and for all contained assertions.
func encodeShard(s *rainslib.ShardSection, context, subjectZone, indent string, addZoneAndContext bool) string {
	rangeFrom := s.RangeFrom
	rangeTo := s.RangeTo
	if rangeFrom == "" {
		rangeFrom = "<"
	}
	if rangeTo == "" {
		rangeTo = ">"
	}
	var shard string
	if addZoneAndContext {
		shard = fmt.Sprintf("%s %s %s %s %s [\n", TypeShard, subjectZone, context, rangeFrom, rangeTo)
	} else {
		shard = fmt.Sprintf("%s %s %s [\n", TypeShard, rangeFrom, rangeTo)
	}
	for _, assertion := range s.Content {
		shard += encodeAssertion(assertion, context, subjectZone, indent+indent4, addZoneAndContext)
	}
	if s.Signatures != nil {
		var sigs []string
		for _, sig := range s.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			return fmt.Sprintf("%s%s%s] ( %s )\n", indent, shard, indent, sigs[0])
		}
		return fmt.Sprintf("%s%s%s] ( \n%s%s\n%s  )\n", indent, shard, indent, indent+indent4, strings.Join(sigs, "\n"+indent+indent4), indent)
	}
	return fmt.Sprintf("%s%s%s]\n", indent, shard, indent)
}

//encodeAssertion returns a in zonefile format. If addZoneAndContext is true, the context and
//subject zone are also present.
func encodeAssertion(a *rainslib.AssertionSection, context, zone, indent string, addZoneAndContext bool) string {
	var assertion string
	if addZoneAndContext {
		assertion = fmt.Sprintf("%s%s %s %s %s [ ", indent, TypeAssertion, a.SubjectName, zone, context)
	} else {
		assertion = fmt.Sprintf("%s%s %s [ ", indent, TypeAssertion, a.SubjectName)
	}
	signature := ""
	if a.Signatures != nil {
		var sigs []string
		for _, sig := range a.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			signature = fmt.Sprintf(" ( %s )\n", sigs[0])
		} else {
			signature = fmt.Sprintf(" ( \n%s%s\n%s)\n", indent+indent4, strings.Join(sigs, "\n"+indent+indent4), indent)
		}
	}
	if len(a.Content) > 1 {
		return fmt.Sprintf("%s\n%s\n%s]%s", assertion, encodeObjects(a.Content, indent+indent4), indent, signature)
	}
	return fmt.Sprintf("%s%s ]%s", assertion, encodeObjects(a.Content, ""), signature)
}

//encodeObjects returns o in zonefile format.
func encodeObjects(o []rainslib.Object, indent string) string {
	var objects []string
	for _, obj := range o {
		object := indent
		switch obj.Type {
		case rainslib.OTName:
			if nameObj, ok := obj.Value.(rainslib.NameObject); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeName), encodeNameObject(nameObj))
			} else {
				log.Error("Type assertion failed. Expected rainslib.NameObject", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTIP6Addr:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeIP6), obj.Value)
		case rainslib.OTIP4Addr:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeIP4), obj.Value)
		case rainslib.OTRedirection:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeRedirection), obj.Value)
		case rainslib.OTDelegation:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeDelegation), encodeEd25519PublicKey(pkey))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTNameset:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeNameSet), obj.Value)
		case rainslib.OTCertInfo:
			if cert, ok := obj.Value.(rainslib.CertificateObject); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeCertificate), encodeCertificate(cert))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.CertificateObject", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTServiceInfo:
			if srvInfo, ok := obj.Value.(rainslib.ServiceInfo); ok {
				object += fmt.Sprintf("%s%s %d %d", addIndentToType(TypeServiceInfo), srvInfo.Name, srvInfo.Port, srvInfo.Priority)
			} else {
				log.Warn("Type assertion failed. Expected rainslib.ServiceInfo", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTRegistrar:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeRegistrar), obj.Value)
		case rainslib.OTRegistrant:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeRegistrant), obj.Value)
		case rainslib.OTInfraKey:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeInfraKey), encodeEd25519PublicKey(pkey))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTExtraKey:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s %s", addIndentToType(TypeExternalKey), encodeKeySpace(pkey.KeySpace), encodeEd25519PublicKey(pkey))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTNextKey:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s %d %d", addIndentToType(TypeNextKey), encodeEd25519PublicKey(pkey), pkey.ValidSince, pkey.ValidUntil)
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		default:
			log.Warn("Unsupported obj type", "type", fmt.Sprintf("%T", obj.Type))
			return ""
		}
		objects = append(objects, object)
	}
	return strings.Join(objects, "\n")
}

//addIndentToType returns the object type ot with appropriate indent such that the object value start at the same
//indent.
func addIndentToType(ot string) string {
	result := "          "
	return ot + result[len(ot):]
}

//encodeNameObject returns no represented as a string in zone file format.
func encodeNameObject(no rainslib.NameObject) string {
	nameObject := []string{}
	for _, oType := range no.Types {
		switch oType {
		case rainslib.OTName:
			nameObject = append(nameObject, TypeName)
		case rainslib.OTIP6Addr:
			nameObject = append(nameObject, TypeIP6)
		case rainslib.OTIP4Addr:
			nameObject = append(nameObject, TypeIP4)
		case rainslib.OTRedirection:
			nameObject = append(nameObject, TypeRedirection)
		case rainslib.OTDelegation:
			nameObject = append(nameObject, TypeDelegation)
		case rainslib.OTNameset:
			nameObject = append(nameObject, TypeNameSet)
		case rainslib.OTCertInfo:
			nameObject = append(nameObject, TypeCertificate)
		case rainslib.OTServiceInfo:
			nameObject = append(nameObject, TypeServiceInfo)
		case rainslib.OTRegistrar:
			nameObject = append(nameObject, TypeRegistrar)
		case rainslib.OTRegistrant:
			nameObject = append(nameObject, TypeRegistrant)
		case rainslib.OTInfraKey:
			nameObject = append(nameObject, TypeInfraKey)
		case rainslib.OTExtraKey:
			nameObject = append(nameObject, TypeExternalKey)
		case rainslib.OTNextKey:
			nameObject = append(nameObject, TypeNextKey)
		default:
			log.Warn("Unsupported object type in nameObject", "actualType", oType, "nameObject", no)
		}
	}
	return fmt.Sprintf("%s [ %s ]", no.Name, strings.Join(nameObject, " "))
}

//encodeEd25519PublicKey returns pkey represented as a string in zone file format.
func encodeEd25519PublicKey(pkey rainslib.PublicKey) string {
	if key, ok := pkey.Key.(ed25519.PublicKey); ok {
		return fmt.Sprintf("%s %d %s", TypeEd25519, pkey.KeyPhase, hex.EncodeToString(key))
	}
	log.Warn("Type assertion failed. Expected rainslib.Ed25519PublicKey", "actualType", fmt.Sprintf("%T", pkey.Key))
	return ""
}

//encodeKeySpace returns keySpace represented as a string in zone file format.
func encodeKeySpace(keySpace rainslib.KeySpaceID) string {
	switch keySpace {
	case rainslib.RainsKeySpace:
		return TypeKSRains
	default:
		log.Warn("Unsupported key space type", "actualType", keySpace)
	}
	return ""
}

//encodeCertificate returns cert represented as a string in zone file format.
func encodeCertificate(cert rainslib.CertificateObject) string {
	var pt, cu, ca string
	switch cert.Type {
	case rainslib.PTUnspecified:
		pt = TypeUnspecified
	case rainslib.PTTLS:
		pt = TypePTTLS
	default:
		log.Warn("Unsupported protocol type", "protocolType", cert.Type)
		return ""
	}
	switch cert.Usage {
	case rainslib.CUTrustAnchor:
		cu = TypeCUTrustAnchor
	case rainslib.CUEndEntity:
		cu = TypeCUEndEntity
	default:
		log.Warn("Unsupported certificate usage", "certUsage", cert.Usage)
		return ""
	}
	switch cert.HashAlgo {
	case rainslib.NoHashAlgo:
		ca = TypeNoHash
	case rainslib.Sha256:
		ca = TypeSha256
	case rainslib.Sha384:
		ca = TypeSha384
	case rainslib.Sha512:
		ca = TypeSha512
	default:
		log.Warn("Unsupported certificate hash algorithm type", "hashAlgoType", cert.HashAlgo)
		return ""
	}
	return fmt.Sprintf("%s %s %s %s", pt, cu, ca, hex.EncodeToString(cert.Data))
}

func encodeEd25519Signature(sig rainslib.Signature) string {
	signature := fmt.Sprintf("%s %s %s %d %d %d", TypeSignature, TypeEd25519, TypeKSRains, sig.PublicKeyID.KeyPhase, sig.ValidSince, sig.ValidUntil)
	if pkey, ok := sig.Data.(ed25519.PublicKey); ok {
		return fmt.Sprintf("%s %s", signature, hex.EncodeToString(pkey))
	}
	return signature
}
<commit_msg>First part pshard encoder<commit_after>package zoneFileParser

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

	"github.com/netsec-ethz/rains/rainslib"

	log "github.com/inconshreveable/log15"
	"golang.org/x/crypto/ed25519"
)

//encodeZone return z in zonefile format. If addZoneAndContext is true, the context and subject zone
//are present also for all contained sections.
func encodeZone(z *rainslib.ZoneSection, addZoneAndContext bool) string {
	zone := fmt.Sprintf("%s %s %s [\n", TypeZone, z.SubjectZone, z.Context)
	for _, section := range z.Content {
		switch section := section.(type) {
		case *rainslib.AssertionSection:
			zone += encodeAssertion(section, z.Context, z.SubjectZone, indent4, addZoneAndContext)
		case *rainslib.ShardSection:
			zone += encodeShard(section, z.Context, z.SubjectZone, indent4, addZoneAndContext)
		default:
			log.Warn("Unsupported message section type", "msgSection", section)
		}
	}
	if z.Signatures != nil {
		var sigs []string
		for _, sig := range z.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			return fmt.Sprintf("%s] ( %s )\n", zone, sigs[0])
		}
		return fmt.Sprintf("%s] ( \n%s%s\n  )\n", zone, indent4, strings.Join(sigs, "\n"+indent4))
	}
	return fmt.Sprintf("%s]\n", zone)
}

//encodeShard returns s in zonefile format. If addZoneAndContext is true, the context and subject
//zone are present for the shard and for all contained assertions.
func encodeShard(s *rainslib.ShardSection, context, subjectZone, indent string, addZoneAndContext bool) string {
	rangeFrom := s.RangeFrom
	rangeTo := s.RangeTo
	if rangeFrom == "" {
		rangeFrom = "<"
	}
	if rangeTo == "" {
		rangeTo = ">"
	}
	var shard string
	if addZoneAndContext {
		shard = fmt.Sprintf("%s %s %s %s %s [\n", TypeShard, subjectZone, context, rangeFrom, rangeTo)
	} else {
		shard = fmt.Sprintf("%s %s %s [\n", TypeShard, rangeFrom, rangeTo)
	}
	for _, assertion := range s.Content {
		shard += encodeAssertion(assertion, context, subjectZone, indent+indent4, addZoneAndContext)
	}
	if s.Signatures != nil {
		var sigs []string
		for _, sig := range s.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			return fmt.Sprintf("%s%s%s] ( %s )\n", indent, shard, indent, sigs[0])
		}
		return fmt.Sprintf("%s%s%s] ( \n%s%s\n%s  )\n", indent, shard, indent, indent+indent4, strings.Join(sigs, "\n"+indent+indent4), indent)
	}
	return fmt.Sprintf("%s%s%s]\n", indent, shard, indent)
}

//encodePshard returns s in zonefile format. If addZoneAndContext is true, the context and subject
//zone are present for the pshard.
func encodePshard(s *rainslib.PshardSection, context, subjectZone, indent string, addZoneAndContext bool) string {
	rangeFrom := s.RangeFrom
	rangeTo := s.RangeTo
	if rangeFrom == "" {
		rangeFrom = "<"
	}
	if rangeTo == "" {
		rangeTo = ">"
	}
	var pshard string
	if addZoneAndContext {
		pshard = fmt.Sprintf("%s %s %s %s %s %s", TypePshard, subjectZone, context, rangeFrom, rangeTo)
	} else {
		pshard = fmt.Sprintf("%s %s %s %s", TypePshard, rangeFrom, rangeTo)
	}
	if s.Signatures != nil {
		var sigs []string
		for _, sig := range s.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			return fmt.Sprintf("%s%s ( %s )\n", indent, pshard, sigs[0])
		}
		return fmt.Sprintf("%s%s ( \n%s%s\n%s  )\n", indent, pshard, indent+indent4, strings.Join(sigs, "\n"+indent+indent4), indent)
	}
	return fmt.Sprintf("%s%s\n", indent, pshard)
}

//encodeDataStructure returns d in zonefile format.
func encodeDataStructure(d *rainslib.DataStructure) string {
	bloomFilter, ok := d.Data.(rainslib.BloomFilter)
	if !ok {
		log.Error("Type not yet implemented", "type", fmt.Sprintf("%T", d.Data))
	}
	var hashFamily []string
	for _, hash := range bloomFilter.HashFamily {
		hashFamily = append(hashFamily, encodeHashAlgo(hash))
	}
	return fmt.Sprintf("%s [ %s ] %d %s %s", d.Type, strings.Join(hashFamily, " "), bloomFilter.NofHashFunctions, d.Type, string(bloomFilter.Filter))
}

//encodeAssertion returns a in zonefile format. If addZoneAndContext is true, the context and
//subject zone are also present.
func encodeAssertion(a *rainslib.AssertionSection, context, zone, indent string, addZoneAndContext bool) string {
	var assertion string
	if addZoneAndContext {
		assertion = fmt.Sprintf("%s%s %s %s %s [ ", indent, TypeAssertion, a.SubjectName, zone, context)
	} else {
		assertion = fmt.Sprintf("%s%s %s [ ", indent, TypeAssertion, a.SubjectName)
	}
	signature := ""
	if a.Signatures != nil {
		var sigs []string
		for _, sig := range a.Signatures {
			sigs = append(sigs, encodeEd25519Signature(sig))
		}
		if len(sigs) == 1 {
			signature = fmt.Sprintf(" ( %s )\n", sigs[0])
		} else {
			signature = fmt.Sprintf(" ( \n%s%s\n%s)\n", indent+indent4, strings.Join(sigs, "\n"+indent+indent4), indent)
		}
	}
	if len(a.Content) > 1 {
		return fmt.Sprintf("%s\n%s\n%s]%s", assertion, encodeObjects(a.Content, indent+indent4), indent, signature)
	}
	return fmt.Sprintf("%s%s ]%s", assertion, encodeObjects(a.Content, ""), signature)
}

//encodeObjects returns o in zonefile format.
func encodeObjects(o []rainslib.Object, indent string) string {
	var objects []string
	for _, obj := range o {
		object := indent
		switch obj.Type {
		case rainslib.OTName:
			if nameObj, ok := obj.Value.(rainslib.NameObject); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeName), encodeNameObject(nameObj))
			} else {
				log.Error("Type assertion failed. Expected rainslib.NameObject", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTIP6Addr:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeIP6), obj.Value)
		case rainslib.OTIP4Addr:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeIP4), obj.Value)
		case rainslib.OTRedirection:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeRedirection), obj.Value)
		case rainslib.OTDelegation:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeDelegation), encodeEd25519PublicKey(pkey))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTNameset:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeNameSet), obj.Value)
		case rainslib.OTCertInfo:
			if cert, ok := obj.Value.(rainslib.CertificateObject); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeCertificate), encodeCertificate(cert))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.CertificateObject", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTServiceInfo:
			if srvInfo, ok := obj.Value.(rainslib.ServiceInfo); ok {
				object += fmt.Sprintf("%s%s %d %d", addIndentToType(TypeServiceInfo), srvInfo.Name, srvInfo.Port, srvInfo.Priority)
			} else {
				log.Warn("Type assertion failed. Expected rainslib.ServiceInfo", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTRegistrar:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeRegistrar), obj.Value)
		case rainslib.OTRegistrant:
			object += fmt.Sprintf("%s%s", addIndentToType(TypeRegistrant), obj.Value)
		case rainslib.OTInfraKey:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s", addIndentToType(TypeInfraKey), encodeEd25519PublicKey(pkey))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTExtraKey:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s %s", addIndentToType(TypeExternalKey), encodeKeySpace(pkey.KeySpace), encodeEd25519PublicKey(pkey))
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		case rainslib.OTNextKey:
			if pkey, ok := obj.Value.(rainslib.PublicKey); ok {
				object += fmt.Sprintf("%s%s %d %d", addIndentToType(TypeNextKey), encodeEd25519PublicKey(pkey), pkey.ValidSince, pkey.ValidUntil)
			} else {
				log.Warn("Type assertion failed. Expected rainslib.PublicKey", "actualType", fmt.Sprintf("%T", obj.Value))
				return ""
			}
		default:
			log.Warn("Unsupported obj type", "type", fmt.Sprintf("%T", obj.Type))
			return ""
		}
		objects = append(objects, object)
	}
	return strings.Join(objects, "\n")
}

//addIndentToType returns the object type ot with appropriate indent such that the object value start at the same
//indent.
func addIndentToType(ot string) string {
	result := "          "
	return ot + result[len(ot):]
}

//encodeNameObject returns no represented as a string in zone file format.
func encodeNameObject(no rainslib.NameObject) string {
	nameObject := []string{}
	for _, oType := range no.Types {
		switch oType {
		case rainslib.OTName:
			nameObject = append(nameObject, TypeName)
		case rainslib.OTIP6Addr:
			nameObject = append(nameObject, TypeIP6)
		case rainslib.OTIP4Addr:
			nameObject = append(nameObject, TypeIP4)
		case rainslib.OTRedirection:
			nameObject = append(nameObject, TypeRedirection)
		case rainslib.OTDelegation:
			nameObject = append(nameObject, TypeDelegation)
		case rainslib.OTNameset:
			nameObject = append(nameObject, TypeNameSet)
		case rainslib.OTCertInfo:
			nameObject = append(nameObject, TypeCertificate)
		case rainslib.OTServiceInfo:
			nameObject = append(nameObject, TypeServiceInfo)
		case rainslib.OTRegistrar:
			nameObject = append(nameObject, TypeRegistrar)
		case rainslib.OTRegistrant:
			nameObject = append(nameObject, TypeRegistrant)
		case rainslib.OTInfraKey:
			nameObject = append(nameObject, TypeInfraKey)
		case rainslib.OTExtraKey:
			nameObject = append(nameObject, TypeExternalKey)
		case rainslib.OTNextKey:
			nameObject = append(nameObject, TypeNextKey)
		default:
			log.Warn("Unsupported object type in nameObject", "actualType", oType, "nameObject", no)
		}
	}
	return fmt.Sprintf("%s [ %s ]", no.Name, strings.Join(nameObject, " "))
}

//encodeEd25519PublicKey returns pkey represented as a string in zone file format.
func encodeEd25519PublicKey(pkey rainslib.PublicKey) string {
	if key, ok := pkey.Key.(ed25519.PublicKey); ok {
		return fmt.Sprintf("%s %d %s", TypeEd25519, pkey.KeyPhase, hex.EncodeToString(key))
	}
	log.Warn("Type assertion failed. Expected rainslib.Ed25519PublicKey", "actualType", fmt.Sprintf("%T", pkey.Key))
	return ""
}

//encodeKeySpace returns keySpace represented as a string in zone file format.
func encodeKeySpace(keySpace rainslib.KeySpaceID) string {
	switch keySpace {
	case rainslib.RainsKeySpace:
		return TypeKSRains
	default:
		log.Warn("Unsupported key space type", "actualType", keySpace)
	}
	return ""
}

//encodeCertificate returns cert represented as a string in zone file format.
func encodeCertificate(cert rainslib.CertificateObject) string {
	var pt, cu, ca string
	switch cert.Type {
	case rainslib.PTUnspecified:
		pt = TypeUnspecified
	case rainslib.PTTLS:
		pt = TypePTTLS
	default:
		log.Warn("Unsupported protocol type", "protocolType", cert.Type)
		return ""
	}
	switch cert.Usage {
	case rainslib.CUTrustAnchor:
		cu = TypeCUTrustAnchor
	case rainslib.CUEndEntity:
		cu = TypeCUEndEntity
	default:
		log.Warn("Unsupported certificate usage", "certUsage", cert.Usage)
		return ""
	}
	ca = encodeHashAlgo(cert.HashAlgo)
	return fmt.Sprintf("%s %s %s %s", pt, cu, ca, hex.EncodeToString(cert.Data))
}

func encodeHashAlgo(h rainslib.HashAlgorithmType) string {
	switch h {
	case rainslib.NoHashAlgo:
		return TypeNoHash
	case rainslib.Sha256:
		return TypeSha256
	case rainslib.Sha384:
		return TypeSha384
	case rainslib.Sha512:
		return TypeSha512
	case rainslib.Fnv64:
		return TypeFnv64
	case rainslib.Murmur364:
		return TypeMurmur364
	default:
		log.Warn("Unsupported certificate hash algorithm type", "hashAlgoType", h)
		return ""
	}
}

func encodeEd25519Signature(sig rainslib.Signature) string {
	signature := fmt.Sprintf("%s %s %s %d %d %d", TypeSignature, TypeEd25519, TypeKSRains, sig.PublicKeyID.KeyPhase, sig.ValidSince, sig.ValidUntil)
	if pkey, ok := sig.Data.(ed25519.PublicKey); ok {
		return fmt.Sprintf("%s %s", signature, hex.EncodeToString(pkey))
	}
	return signature
}
<|endoftext|>


<commit_before>package mcquery

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"strconv"
	"time"
)

const (
	handshakeType = 9
	statusType    = 0
)

func magicBytes() [2]byte {
	return [2]byte{0xFE, 0xFD}
}

type handshakeRequest struct {
	Magic     [2]byte
	Type      byte
	SessionId int32
}

type basicStatRequest struct {
	Magic          [2]byte
	Type           byte
	SessionId      int32
	ChallengeToken int32
}

type responseHeader struct {
	Type      byte
	SessionId int32
}

type BasicStatResponse struct {
	Type       byte
	SessionId  int32
	Motd       string
	Gametype   string
	Map        string
	NumPlayers string
	MaxPlayers string
	HostPort   uint16
	HostIp     string
}

func Connect(ip string, port uint16) (*bufio.ReadWriter, error, chan<- bool) {
	conn, err := net.Dial("udp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err, nil
	}

	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

	kill := make(chan bool)
	go func() {
		<-kill
		conn.Close()
	}()

	return rw, nil, kill
}

func BasicStat(rw *bufio.ReadWriter, challenge int32) (*BasicStatResponse, error) {

	var response BasicStatResponse
	done := make(chan error)
	request := basicStatRequest{
		Magic:          magicBytes(),
		Type:           statusType,
		SessionId:      1,
		ChallengeToken: challenge,
	}

	go func() {

		err := binary.Write(rw, binary.BigEndian, request)
		if err != nil {
			done <- err
			return
		}
		rw.Flush()

		var responseHeader responseHeader
		err = binary.Read(rw, binary.BigEndian, &responseHeader)
		if err != nil {
			done <- err
			return
		}

		// All of this code here is really ugly, makes me want to switch to
		// haskell and use monads...

		mesg, err := rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]

		response.Motd = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.Gametype = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.Map = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.NumPlayers = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.MaxPlayers = string(mesg)

		var wrappedPort struct {
			Port uint16
		}
		err = binary.Read(rw, binary.LittleEndian, &wrappedPort)
		if err != nil {
			done <- err
			return
		}
		response.HostPort = wrappedPort.Port

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.HostIp = string(mesg)

		done <- nil
	}()

	select {
	case <-time.After(2000 * time.Millisecond):
		return nil, errors.New("Timed out")
	case err := <-done:
		if err == nil {
			return &response, nil
		} else {
			return nil, err
		}
	}
}

// Handshakes with the minecraft server, returning the challenge token
func Handshake(rw *bufio.ReadWriter) (int32, error) {

	var response responseHeader
	var challenge int32
	done := make(chan error)
	request := handshakeRequest{
		Magic:     magicBytes(),
		Type:      handshakeType,
		SessionId: 1, // TODO figure this out
	}

	go func() {
		err := binary.Write(rw, binary.BigEndian, request)
		if err != nil {
			done <- err
			return
		}
		if err = rw.Flush(); err != nil {
			done <- err
			return
		}
		if err = binary.Read(rw, binary.BigEndian, &response); err != nil {
			done <- err
			return
		}

		payload, err := rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		payload = payload[:len(payload)-1]

		number, err := strconv.ParseInt(string(payload), 10, 32)
		if err != nil {
			done <- err
			return
		}
		challenge = int32(number)
		done <- nil
	}()

	select {
	case <-time.After(2 * time.Second):
		return 0, errors.New("Timed out")
	case err := <-done:
		if err == nil {
			return challenge, nil
		} else {
			return 0, err
		}
	}

}
<commit_msg>Added full stat query<commit_after>package mcquery

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"strconv"
	"time"
)

const (
	handshakeType = 9
	statusType    = 0
)

// Max amount of times looped readers can loop for, in case key values
// or player list is malformed
const loopLimit = 512

func magicBytes() [2]byte {
	return [2]byte{0xFE, 0xFD}
}

type handshakeRequest struct {
	Magic     [2]byte
	Type      byte
	SessionId int32
}

type basicStatRequest struct {
	Magic          [2]byte
	Type           byte
	SessionId      int32
	ChallengeToken int32
}

type fullStatRequest struct {
	Magic          [2]byte
	Type           byte
	SessionId      int32
	ChallengeToken int32
	Padding        [4]byte
}

type FullStatResponse struct {
	Type      byte
	SessionId int32
	Padding   [11]byte
	KeyValues map[string]string
	Padding2  [10]byte
	Players   []string
}

type responseHeader struct {
	Type      byte
	SessionId int32
}

type BasicStatResponse struct {
	Type       byte
	SessionId  int32
	Motd       string
	Gametype   string
	Map        string
	NumPlayers string
	MaxPlayers string
	HostPort   uint16
	HostIp     string
}

func Connect(ip string, port uint16) (*bufio.ReadWriter, error, chan<- bool) {
	conn, err := net.Dial("udp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err, nil
	}

	rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

	kill := make(chan bool)
	go func() {
		<-kill
		conn.Close()
	}()

	return rw, nil, kill
}

func FullStat(rw *bufio.ReadWriter, challenge int32) (*FullStatResponse, error) {

	var response FullStatResponse
	done := make(chan error)
	request := fullStatRequest{
		Magic:          magicBytes(),
		Type:           statusType,
		SessionId:      1,
		ChallengeToken: challenge,
	}

	go func() {

		err := binary.Write(rw, binary.BigEndian, request)
		if err != nil {
			done <- err
			return
		}
		rw.Flush()

		var header responseHeader
		err = binary.Read(rw, binary.BigEndian, &header)
		if err != nil {
			done <- err
			return
		}
		response.Type = header.Type
		response.SessionId = header.SessionId

		_, err = rw.Discard(11)
		if err != nil {
			done <- err
			return
		}

		response.KeyValues = make(map[string]string)
		for i := 0; ; i += 1 {
			var key, value []byte
			key, err = rw.ReadBytes(0)
			if err != nil {
				done <- err
				return
			}
			key = key[:len(key)-1]
			if len(key) == 0 {
				break
			}
			value, err = rw.ReadBytes(0)
			if err != nil {
				done <- err
				return
			}
			value = value[:len(value)-1]
			response.KeyValues[string(key)] = string(value)
			if i > loopLimit {
				done <- errors.New("Too many key values!")
				return
			}
		}

		_, err = rw.Discard(10)
		if err != nil {
			done <- err
			return
		}

		response.Players = make([]string, 0)
		for i := 0; ; i += 1 {

			playerName, err := rw.ReadBytes(0)
			if err != nil {
				done <- err
				return
			}
			playerName = playerName[:len(playerName)-1]
			if len(playerName) == 0 {
				break
			}
			response.Players = append(response.Players, string(playerName))

			if i > loopLimit {
				done <- errors.New("Too many players!")
				return
			}
		}

		done <- nil
	}()

	select {
	case <-time.After(2 * time.Second):
		return nil, nil
	case err := <-done:
		if err == nil {
			return &response, nil
		} else {
			return nil, err
		}
	}
}

func BasicStat(rw *bufio.ReadWriter, challenge int32) (*BasicStatResponse, error) {

	var response BasicStatResponse
	done := make(chan error)
	request := basicStatRequest{
		Magic:          magicBytes(),
		Type:           statusType,
		SessionId:      1,
		ChallengeToken: challenge,
	}

	go func() {

		err := binary.Write(rw, binary.BigEndian, request)
		if err != nil {
			done <- err
			return
		}
		rw.Flush()

		var responseHeader responseHeader
		err = binary.Read(rw, binary.BigEndian, &responseHeader)
		if err != nil {
			done <- err
			return
		}
		response.Type = responseHeader.Type
		response.SessionId = responseHeader.SessionId

		// All of this code here is really ugly, makes me want to switch to
		// haskell and use monads...

		mesg, err := rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]

		response.Motd = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.Gametype = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.Map = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.NumPlayers = string(mesg)

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.MaxPlayers = string(mesg)

		var wrappedPort struct {
			Port uint16
		}
		err = binary.Read(rw, binary.LittleEndian, &wrappedPort)
		if err != nil {
			done <- err
			return
		}
		response.HostPort = wrappedPort.Port

		mesg, err = rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		mesg = mesg[:len(mesg)-1]
		response.HostIp = string(mesg)

		done <- nil
	}()

	select {
	case <-time.After(2000 * time.Millisecond):
		return nil, errors.New("Timed out")
	case err := <-done:
		if err == nil {
			return &response, nil
		} else {
			return nil, err
		}
	}
}

// Handshakes with the minecraft server, returning the challenge token
func Handshake(rw *bufio.ReadWriter) (int32, error) {

	var response responseHeader
	var challenge int32
	done := make(chan error)
	request := handshakeRequest{
		Magic:     magicBytes(),
		Type:      handshakeType,
		SessionId: 1, // TODO figure this out
	}

	go func() {
		err := binary.Write(rw, binary.BigEndian, request)
		if err != nil {
			done <- err
			return
		}
		if err = rw.Flush(); err != nil {
			done <- err
			return
		}
		if err = binary.Read(rw, binary.BigEndian, &response); err != nil {
			done <- err
			return
		}

		payload, err := rw.ReadBytes(0)
		if err != nil {
			done <- err
			return
		}
		payload = payload[:len(payload)-1]

		number, err := strconv.ParseInt(string(payload), 10, 32)
		if err != nil {
			done <- err
			return
		}
		challenge = int32(number)
		done <- nil
	}()

	select {
	case <-time.After(2 * time.Second):
		return 0, errors.New("Timed out")
	case err := <-done:
		if err == nil {
			return challenge, nil
		} else {
			return 0, err
		}
	}

}
<|endoftext|>