<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 blob

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"log"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
)

// A key placed in GCS object metadata by GCSStore containing the hex SHA-1
// expected for the object contents. This is of course redundant with the
// object name; we use it as a paranoid check against GCS returning the
// metadata or contents for the wrong object.
const metadataKey_SHA1 = "comeback_sha1"

// A key placed in GCS object metadata by GCSStore containing the CRC32C
// checksum expected for the object contents. If GCS reports a different
// checksum or returns contents with a different checksum, we know something
// screwy has happened.
//
// See here for more info: https://github.com/jacobsa/comeback/issues/18
const metadataKey_CRC32C = "comeback_crc32c"

// A key placed in GCS object metadata by GCSStore containing the hex MD5 sum
// expected for the object contents. If GCS reports a different MD5 sum or
// returns contents with a different MD5 sum, we know something screwy has
// happened.
//
// See here for more info: https://github.com/jacobsa/comeback/issues/18
const metadataKey_MD5 = "comeback_md5"

// Return a blob store that stores blobs in the supplied GCS bucket. GCS object
// names look like:
//
//     <prefix><score>
//
// where <score> is the result of calling Score.Hex.
//
// The blob store trusts that it has full ownership of this portion of the
// bucket's namespace -- if a score name exists, then it points to the correct
// data.
//
// The returned store does not support Flush or Contains; these methods must
// not be called.
func NewGCSStore(
	bucket gcs.Bucket,
	prefix string) (store *GCSStore) {
	store = &GCSStore{
		bucket:     bucket,
		namePrefix: prefix,
	}

	return
}

type GCSStore struct {
	bucket     gcs.Bucket
	namePrefix string
}

// Parse and verify the internal consistency of the supplied object record in
// the same manner that a GCSStore configured with the supplied object name
// prefix would. Return the score of the blob that the object contains.
func ParseObjectRecord(
	o *gcs.Object,
	namePrefix string) (score Score, err error) {
	// Is the name of the appropriate form?
	if !strings.HasPrefix(o.Name, namePrefix) {
		err = fmt.Errorf("Unexpected object name: %q", o.Name)
		return
	}

	// Parse the hex score.
	hexScore := strings.TrimPrefix(o.Name, namePrefix)
	score, err = ParseHexScore(hexScore)
	if err != nil {
		err = fmt.Errorf(
			"Unexpected hex score %q for object %q: %v",
			hexScore,
			o.Name,
			err)
		return
	}

	// We expect the hex score to match the hex SHA-1 in the metadata.
	hexSHA1, ok := o.Metadata[metadataKey_SHA1]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_SHA1)
		return
	}

	if hexSHA1 != hexScore {
		err = fmt.Errorf(
			"Score/SHA-1 metadata mismatch for object %q: %q",
			o.Name,
			hexSHA1)
		return
	}

	// We expect the hex CRC32C in the object metadata to match what GCS says the
	// object's checksum is.
	hexCRC32C, ok := o.Metadata[metadataKey_CRC32C]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_CRC32C)
		return
	}

	crc32Uint64, err := strconv.ParseUint(hexCRC32C, 0, 32)
	if err != nil {
		err = fmt.Errorf(
			"Object %q has invalid hex CRC32C %q: %v",
			o.Name,
			hexCRC32C,
			err)
		return
	}

	if uint32(crc32Uint64) != o.CRC32C {
		err = fmt.Errorf(
			"CRC32C mismatch for object %q: %#08x vs. %#08x",
			o.Name,
			crc32Uint64,
			o.CRC32C)
		return
	}

	// We expect the hex MD5 in the object metadata to match what GCS says the
	// object's MD5 sum is.
	hexMD5, ok := o.Metadata[metadataKey_MD5]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_MD5)
		return
	}

	if hex.DecodedLen(len(hexMD5)) != md5.Size {
		err = fmt.Errorf(
			"Object %q has weird hex MD5 metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	var md5 [md5.Size]byte
	_, err = hex.Decode(md5[:], []byte(hexMD5))
	if err != nil {
		err = fmt.Errorf(
			"Object %q has invalid hex MD5 in metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	if md5 != o.MD5 {
		err = fmt.Errorf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(md5[:]),
			hex.EncodeToString(o.MD5[:]))
		return
	}

	return
}

// Write object records for all of the blob objects in the supplied bucket into
// the given channel, without closing it. The order of records is undefined.
// The caller will likely want to call ParseObjectRecord for each record.
func ListBlobObjects(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string,
	objects chan<- *gcs.Object) (err error) {
	req := &gcs.ListObjectsRequest{
		Prefix:            blobObjectNamePrefix,
		ContinuationToken: *fToken,
	}

	// List until we run out.
	for {
		// Fetch the next batch.
		var listing *gcs.Listing
		listing, err = bucket.ListObjects(ctx, req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Pass on each object.
		for _, o := range listing.Objects {
			// Special case: for gcsfuse compatibility, we allow blobObjectNamePrefix
			// to exist as its own object name. Skip it.
			if o.Name == blobObjectNamePrefix {
				continue
			}

			select {
			case objects <- o:

				// Cancelled?
			case <-ctx.Done():
				err = ctx.Err()
				return
			}
		}

		// Are we done?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
		log.Printf("Continuation token: %q", req.ContinuationToken)
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) makeName(score Score) (name string) {
	name = s.namePrefix + score.Hex()
	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) Store(blob []byte) (score Score, err error) {
	// Compute a score and an object name.
	score = ComputeScore(blob)
	name := s.makeName(score)

	// Create the object.
	crc32c := *gcsutil.CRC32C(blob)
	md5 := *gcsutil.MD5(blob)
	sha1 := sha1.Sum(blob)

	req := &gcs.CreateObjectRequest{
		Name:     name,
		Contents: bytes.NewReader(blob),
		CRC32C:   &crc32c,
		MD5:      &md5,

		Metadata: map[string]string{
			metadataKey_SHA1:   hex.EncodeToString(sha1[:]),
			metadataKey_CRC32C: fmt.Sprintf("%#08x", crc32c),
			metadataKey_MD5:    hex.EncodeToString(md5[:]),
		},
	}

	o, err := s.bucket.CreateObject(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("CreateObject: %v", err)
		return
	}

	// Paranoid check: what we get back from GCS should match what we put in.
	if o.CRC32C != crc32c {
		panic(fmt.Sprintf(
			"CRC32C mismatch for object %q: 0x%08xv vs. 0x%08x",
			o.Name,
			o.CRC32C,
			crc32c))
	}

	if o.MD5 != md5 {
		panic(fmt.Sprintf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(o.MD5[:]),
			hex.EncodeToString(md5[:])))
	}

	return
}

func (s *GCSStore) Flush() (err error) {
	panic("GCSStore.Flush not supported; wiring code bug?")
}

func (s *GCSStore) Contains(score Score) (b bool) {
	panic("GCSStore.Contains not supported; wiring code bug?")
}

func (s *GCSStore) Load(score Score) (blob []byte, err error) {
	// Create a ReadCloser.
	req := &gcs.ReadObjectRequest{
		Name: s.makeName(score),
	}

	rc, err := s.bucket.NewReader(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	// Read from it.
	blob, err = ioutil.ReadAll(rc)
	if err != nil {
		rc.Close()
		err = fmt.Errorf("ReadAll: %v", err)
		return
	}

	// Close it.
	err = rc.Close()
	if err != nil {
		err = fmt.Errorf("Close: %v", err)
		return
	}

	return
}

// List all of the blobs that are known to be durable in the bucket.
func (s *GCSStore) List() (scores []Score, err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: s.namePrefix,
	}

	// List repeatedly until we're done.
	for {
		// Call the bucket.
		var listing *gcs.Listing
		listing, err = s.bucket.ListObjects(context.Background(), req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Process results.
		for _, o := range listing.Objects {
			// Special case: listing "blobs/*" includes "blobs/" itself, which we
			// allow to exist for convenience of use with e.g. gcsfuse.
			if o.Name == s.namePrefix {
				continue
			}

			// Parse and verify the record.
			var score Score
			score, err = ParseObjectRecord(o, s.namePrefix)
			if err != nil {
				return
			}

			// Save the score.
			scores = append(scores, score)
		}

		// Continue?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}
<commit_msg>Fixed build errors.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 blob

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
)

// A key placed in GCS object metadata by GCSStore containing the hex SHA-1
// expected for the object contents. This is of course redundant with the
// object name; we use it as a paranoid check against GCS returning the
// metadata or contents for the wrong object.
const metadataKey_SHA1 = "comeback_sha1"

// A key placed in GCS object metadata by GCSStore containing the CRC32C
// checksum expected for the object contents. If GCS reports a different
// checksum or returns contents with a different checksum, we know something
// screwy has happened.
//
// See here for more info: https://github.com/jacobsa/comeback/issues/18
const metadataKey_CRC32C = "comeback_crc32c"

// A key placed in GCS object metadata by GCSStore containing the hex MD5 sum
// expected for the object contents. If GCS reports a different MD5 sum or
// returns contents with a different MD5 sum, we know something screwy has
// happened.
//
// See here for more info: https://github.com/jacobsa/comeback/issues/18
const metadataKey_MD5 = "comeback_md5"

// Return a blob store that stores blobs in the supplied GCS bucket. GCS object
// names look like:
//
//     <prefix><score>
//
// where <score> is the result of calling Score.Hex.
//
// The blob store trusts that it has full ownership of this portion of the
// bucket's namespace -- if a score name exists, then it points to the correct
// data.
//
// The returned store does not support Flush or Contains; these methods must
// not be called.
func NewGCSStore(
	bucket gcs.Bucket,
	prefix string) (store *GCSStore) {
	store = &GCSStore{
		bucket:     bucket,
		namePrefix: prefix,
	}

	return
}

type GCSStore struct {
	bucket     gcs.Bucket
	namePrefix string
}

// Parse and verify the internal consistency of the supplied object record in
// the same manner that a GCSStore configured with the supplied object name
// prefix would. Return the score of the blob that the object contains.
func ParseObjectRecord(
	o *gcs.Object,
	namePrefix string) (score Score, err error) {
	// Is the name of the appropriate form?
	if !strings.HasPrefix(o.Name, namePrefix) {
		err = fmt.Errorf("Unexpected object name: %q", o.Name)
		return
	}

	// Parse the hex score.
	hexScore := strings.TrimPrefix(o.Name, namePrefix)
	score, err = ParseHexScore(hexScore)
	if err != nil {
		err = fmt.Errorf(
			"Unexpected hex score %q for object %q: %v",
			hexScore,
			o.Name,
			err)
		return
	}

	// We expect the hex score to match the hex SHA-1 in the metadata.
	hexSHA1, ok := o.Metadata[metadataKey_SHA1]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_SHA1)
		return
	}

	if hexSHA1 != hexScore {
		err = fmt.Errorf(
			"Score/SHA-1 metadata mismatch for object %q: %q",
			o.Name,
			hexSHA1)
		return
	}

	// We expect the hex CRC32C in the object metadata to match what GCS says the
	// object's checksum is.
	hexCRC32C, ok := o.Metadata[metadataKey_CRC32C]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_CRC32C)
		return
	}

	crc32Uint64, err := strconv.ParseUint(hexCRC32C, 0, 32)
	if err != nil {
		err = fmt.Errorf(
			"Object %q has invalid hex CRC32C %q: %v",
			o.Name,
			hexCRC32C,
			err)
		return
	}

	if uint32(crc32Uint64) != o.CRC32C {
		err = fmt.Errorf(
			"CRC32C mismatch for object %q: %#08x vs. %#08x",
			o.Name,
			crc32Uint64,
			o.CRC32C)
		return
	}

	// We expect the hex MD5 in the object metadata to match what GCS says the
	// object's MD5 sum is.
	hexMD5, ok := o.Metadata[metadataKey_MD5]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_MD5)
		return
	}

	if hex.DecodedLen(len(hexMD5)) != md5.Size {
		err = fmt.Errorf(
			"Object %q has weird hex MD5 metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	var md5 [md5.Size]byte
	_, err = hex.Decode(md5[:], []byte(hexMD5))
	if err != nil {
		err = fmt.Errorf(
			"Object %q has invalid hex MD5 in metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	if md5 != o.MD5 {
		err = fmt.Errorf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(md5[:]),
			hex.EncodeToString(o.MD5[:]))
		return
	}

	return
}

// Write object records for all of the blob objects in the supplied bucket into
// the given channel, without closing it. The order of records is undefined.
// The caller will likely want to call ParseObjectRecord for each record.
func ListBlobObjects(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string,
	objects chan<- *gcs.Object) (err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: namePrefix,
	}

	// List until we run out.
	for {
		// Fetch the next batch.
		var listing *gcs.Listing
		listing, err = bucket.ListObjects(ctx, req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Pass on each object.
		for _, o := range listing.Objects {
			// Special case: for gcsfuse compatibility, we allow namePrefix to exist
			// as its own object name. Skip it.
			if o.Name == namePrefix {
				continue
			}

			select {
			case objects <- o:

				// Cancelled?
			case <-ctx.Done():
				err = ctx.Err()
				return
			}
		}

		// Are we done?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) makeName(score Score) (name string) {
	name = s.namePrefix + score.Hex()
	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) Store(blob []byte) (score Score, err error) {
	// Compute a score and an object name.
	score = ComputeScore(blob)
	name := s.makeName(score)

	// Create the object.
	crc32c := *gcsutil.CRC32C(blob)
	md5 := *gcsutil.MD5(blob)
	sha1 := sha1.Sum(blob)

	req := &gcs.CreateObjectRequest{
		Name:     name,
		Contents: bytes.NewReader(blob),
		CRC32C:   &crc32c,
		MD5:      &md5,

		Metadata: map[string]string{
			metadataKey_SHA1:   hex.EncodeToString(sha1[:]),
			metadataKey_CRC32C: fmt.Sprintf("%#08x", crc32c),
			metadataKey_MD5:    hex.EncodeToString(md5[:]),
		},
	}

	o, err := s.bucket.CreateObject(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("CreateObject: %v", err)
		return
	}

	// Paranoid check: what we get back from GCS should match what we put in.
	if o.CRC32C != crc32c {
		panic(fmt.Sprintf(
			"CRC32C mismatch for object %q: 0x%08xv vs. 0x%08x",
			o.Name,
			o.CRC32C,
			crc32c))
	}

	if o.MD5 != md5 {
		panic(fmt.Sprintf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(o.MD5[:]),
			hex.EncodeToString(md5[:])))
	}

	return
}

func (s *GCSStore) Flush() (err error) {
	panic("GCSStore.Flush not supported; wiring code bug?")
}

func (s *GCSStore) Contains(score Score) (b bool) {
	panic("GCSStore.Contains not supported; wiring code bug?")
}

func (s *GCSStore) Load(score Score) (blob []byte, err error) {
	// Create a ReadCloser.
	req := &gcs.ReadObjectRequest{
		Name: s.makeName(score),
	}

	rc, err := s.bucket.NewReader(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	// Read from it.
	blob, err = ioutil.ReadAll(rc)
	if err != nil {
		rc.Close()
		err = fmt.Errorf("ReadAll: %v", err)
		return
	}

	// Close it.
	err = rc.Close()
	if err != nil {
		err = fmt.Errorf("Close: %v", err)
		return
	}

	return
}

// List all of the blobs that are known to be durable in the bucket.
func (s *GCSStore) List() (scores []Score, err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: s.namePrefix,
	}

	// List repeatedly until we're done.
	for {
		// Call the bucket.
		var listing *gcs.Listing
		listing, err = s.bucket.ListObjects(context.Background(), req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Process results.
		for _, o := range listing.Objects {
			// Special case: listing "blobs/*" includes "blobs/" itself, which we
			// allow to exist for convenience of use with e.g. gcsfuse.
			if o.Name == s.namePrefix {
				continue
			}

			// Parse and verify the record.
			var score Score
			score, err = ParseObjectRecord(o, s.namePrefix)
			if err != nil {
				return
			}

			// Save the score.
			scores = append(scores, score)
		}

		// Continue?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}
<|endoftext|>


<commit_before>package consul

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/hashicorp/consul/testutil/retry"
	"github.com/hashicorp/go-uuid"
	"github.com/hashicorp/net-rpc-msgpackrpc"
)

func generateUUID() (ret string) {
	var err error
	if ret, err = uuid.GenerateUUID(); err != nil {
		panic(fmt.Sprintf("Unable to generate a UUID, %v", err))
	}
	return ret
}

func TestInitializeSessionTimers(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "10s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Reset the session timers
	err := s1.initializeSessionTimers()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that we have a timer
	if s1.sessionTimers.Get(session.ID) == nil {
		t.Fatalf("missing session timer")
	}
}

func TestResetSessionTimer_Fault(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	// Should not exist
	err := s1.resetSessionTimer(generateUUID(), nil)
	if err == nil || !strings.Contains(err.Error(), "not found") {
		t.Fatalf("err: %v", err)
	}

	// Create a session
	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "10s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Reset the session timer
	err = s1.resetSessionTimer(session.ID, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that we have a timer
	if s1.sessionTimers.Get(session.ID) == nil {
		t.Fatalf("missing session timer")
	}
}

func TestResetSessionTimer_NoTTL(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	// Create a session
	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "0000s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Reset the session timer
	err := s1.resetSessionTimer(session.ID, session)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that we have a timer
	if s1.sessionTimers.Get(session.ID) != nil {
		t.Fatalf("should not have session timer")
	}
}

func TestResetSessionTimer_InvalidTTL(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	// Create a session
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "foo",
	}

	// Reset the session timer
	err := s1.resetSessionTimer(session.ID, session)
	if err == nil || !strings.Contains(err.Error(), "Invalid Session TTL") {
		t.Fatalf("err: %v", err)
	}
}

func TestResetSessionTimerLocked(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	s1.createSessionTimer("foo", 5*time.Millisecond)
	if s1.sessionTimers.Get("foo") == nil {
		t.Fatalf("missing timer")
	}

	time.Sleep(10 * time.Millisecond * structs.SessionTTLMultiplier)
	if s1.sessionTimers.Get("foo") != nil {
		t.Fatalf("timer should be gone")
	}
}

func TestResetSessionTimerLocked_Renew(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	ttl := 100 * time.Millisecond

	// create the timer
	s1.createSessionTimer("foo", ttl)
	if s1.sessionTimers.Get("foo") == nil {
		t.Fatalf("missing timer")
	}

	// wait until it is "expired" but at this point
	// the session still exists.
	time.Sleep(ttl)
	if s1.sessionTimers.Get("foo") == nil {
		t.Fatal("missing timer")
	}

	// renew the session which will reset the TTL to 2*ttl
	// since that is the current SessionTTLMultiplier
	s1.createSessionTimer("foo", ttl)

	// Watch for invalidation
	renew := time.Now()
	deadline := renew.Add(2 * structs.SessionTTLMultiplier * ttl)
	for {
		now := time.Now()
		if now.After(deadline) {
			t.Fatal("should have expired by now")
		}

		// timer still exists
		if s1.sessionTimers.Get("foo") != nil {
			time.Sleep(time.Millisecond)
			continue
		}

		// timer gone
		if now.Sub(renew) < ttl {
			t.Fatalf("early invalidate")
		}
		break
	}
}

func TestInvalidateSession(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	// Create a session
	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "10s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// This should cause a destroy
	s1.invalidateSession(session.ID)

	// Check it is gone
	_, sess, err := state.SessionGet(nil, session.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if sess != nil {
		t.Fatalf("should destroy session")
	}
}

func TestClearSessionTimer(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	s1.createSessionTimer("foo", 5*time.Millisecond)

	err := s1.clearSessionTimer("foo")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if s1.sessionTimers.Get("foo") != nil {
		t.Fatalf("timer should be gone")
	}
}

func TestClearAllSessionTimers(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	s1.createSessionTimer("foo", 10*time.Millisecond)
	s1.createSessionTimer("bar", 10*time.Millisecond)
	s1.createSessionTimer("baz", 10*time.Millisecond)

	err := s1.clearAllSessionTimers()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// sessionTimers is guarded by the lock
	if s1.sessionTimers.Len() != 0 {
		t.Fatalf("timers should be gone")
	}
}

func TestServer_SessionTTL_Failover(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()
	testrpc.WaitForTestAgent(t, s1.RPC, "dc1")

	dir2, s2 := testServerDCBootstrap(t, "dc1", false)
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	dir3, s3 := testServerDCBootstrap(t, "dc1", false)
	defer os.RemoveAll(dir3)
	defer s3.Shutdown()
	servers := []*Server{s1, s2, s3}

	// Try to join
	joinLAN(t, s2, s1)
	joinLAN(t, s3, s1)
	retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s1, 3)) })

	// Find the leader
	var leader *Server
	for _, s := range servers {
		// Check that s.sessionTimers is empty
		if s.sessionTimers.Len() != 0 {
			t.Fatalf("should have no sessionTimers")
		}
		// Find the leader too
		if s.IsLeader() {
			leader = s
		}
	}
	if leader == nil {
		t.Fatalf("Should have a leader")
	}

	codec := rpcClient(t, leader)
	defer codec.Close()

	// Register a node
	node := structs.RegisterRequest{
		Datacenter: s1.config.Datacenter,
		Node:       "foo",
		Address:    "127.0.0.1",
	}
	var out struct{}
	if err := s1.RPC("Catalog.Register", &node, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a TTL session
	arg := structs.SessionRequest{
		Datacenter: "dc1",
		Op:         structs.SessionCreate,
		Session: structs.Session{
			Node: "foo",
			TTL:  "10s",
		},
	}
	var id1 string
	if err := msgpackrpc.CallWithCodec(codec, "Session.Apply", &arg, &id1); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that sessionTimers has the session ID
	if leader.sessionTimers.Get(id1) == nil {
		t.Fatalf("missing session timer")
	}

	// Shutdown the leader!
	leader.Shutdown()

	// sessionTimers should be cleared on leader shutdown
	if leader.sessionTimers.Len() != 0 {
		t.Fatalf("session timers should be empty on the shutdown leader")
	}
	// Find the new leader
	retry.Run(t, func(r *retry.R) {
		leader = nil
		for _, s := range servers {
			if s.IsLeader() {
				leader = s
			}
		}
		if leader == nil {
			r.Fatal("Should have a new leader")
		}

		// Ensure session timer is restored
		if leader.sessionTimers.Get(id1) == nil {
			r.Fatal("missing session timer")
		}
	})
}
<commit_msg>add retry to TestResetSessionTimerLocked<commit_after>package consul

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/hashicorp/consul/testutil/retry"
	"github.com/hashicorp/go-uuid"
	"github.com/hashicorp/net-rpc-msgpackrpc"
)

func generateUUID() (ret string) {
	var err error
	if ret, err = uuid.GenerateUUID(); err != nil {
		panic(fmt.Sprintf("Unable to generate a UUID, %v", err))
	}
	return ret
}

func TestInitializeSessionTimers(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "10s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Reset the session timers
	err := s1.initializeSessionTimers()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that we have a timer
	if s1.sessionTimers.Get(session.ID) == nil {
		t.Fatalf("missing session timer")
	}
}

func TestResetSessionTimer_Fault(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	// Should not exist
	err := s1.resetSessionTimer(generateUUID(), nil)
	if err == nil || !strings.Contains(err.Error(), "not found") {
		t.Fatalf("err: %v", err)
	}

	// Create a session
	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "10s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Reset the session timer
	err = s1.resetSessionTimer(session.ID, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that we have a timer
	if s1.sessionTimers.Get(session.ID) == nil {
		t.Fatalf("missing session timer")
	}
}

func TestResetSessionTimer_NoTTL(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	// Create a session
	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "0000s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Reset the session timer
	err := s1.resetSessionTimer(session.ID, session)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that we have a timer
	if s1.sessionTimers.Get(session.ID) != nil {
		t.Fatalf("should not have session timer")
	}
}

func TestResetSessionTimer_InvalidTTL(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	// Create a session
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "foo",
	}

	// Reset the session timer
	err := s1.resetSessionTimer(session.ID, session)
	if err == nil || !strings.Contains(err.Error(), "Invalid Session TTL") {
		t.Fatalf("err: %v", err)
	}
}

func TestResetSessionTimerLocked(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	s1.createSessionTimer("foo", 5*time.Millisecond)
	if s1.sessionTimers.Get("foo") == nil {
		t.Fatalf("missing timer")
	}

	retry.Run(t, func(r *retry.R) {
		if s1.sessionTimers.Get("foo") != nil {
			r.Fatal("timer should be gone")
		}
	})
}

func TestResetSessionTimerLocked_Renew(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	ttl := 100 * time.Millisecond

	// create the timer
	s1.createSessionTimer("foo", ttl)
	if s1.sessionTimers.Get("foo") == nil {
		t.Fatalf("missing timer")
	}

	// wait until it is "expired" but at this point
	// the session still exists.
	time.Sleep(ttl)
	if s1.sessionTimers.Get("foo") == nil {
		t.Fatal("missing timer")
	}

	// renew the session which will reset the TTL to 2*ttl
	// since that is the current SessionTTLMultiplier
	s1.createSessionTimer("foo", ttl)

	// Watch for invalidation
	renew := time.Now()
	deadline := renew.Add(2 * structs.SessionTTLMultiplier * ttl)
	for {
		now := time.Now()
		if now.After(deadline) {
			t.Fatal("should have expired by now")
		}

		// timer still exists
		if s1.sessionTimers.Get("foo") != nil {
			time.Sleep(time.Millisecond)
			continue
		}

		// timer gone
		if now.Sub(renew) < ttl {
			t.Fatalf("early invalidate")
		}
		break
	}
}

func TestInvalidateSession(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	testrpc.WaitForLeader(t, s1.RPC, "dc1")

	// Create a session
	state := s1.fsm.State()
	if err := state.EnsureNode(1, &structs.Node{Node: "foo", Address: "127.0.0.1"}); err != nil {
		t.Fatalf("err: %s", err)
	}
	session := &structs.Session{
		ID:   generateUUID(),
		Node: "foo",
		TTL:  "10s",
	}
	if err := state.SessionCreate(100, session); err != nil {
		t.Fatalf("err: %v", err)
	}

	// This should cause a destroy
	s1.invalidateSession(session.ID)

	// Check it is gone
	_, sess, err := state.SessionGet(nil, session.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if sess != nil {
		t.Fatalf("should destroy session")
	}
}

func TestClearSessionTimer(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	s1.createSessionTimer("foo", 5*time.Millisecond)

	err := s1.clearSessionTimer("foo")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if s1.sessionTimers.Get("foo") != nil {
		t.Fatalf("timer should be gone")
	}
}

func TestClearAllSessionTimers(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()

	s1.createSessionTimer("foo", 10*time.Millisecond)
	s1.createSessionTimer("bar", 10*time.Millisecond)
	s1.createSessionTimer("baz", 10*time.Millisecond)

	err := s1.clearAllSessionTimers()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// sessionTimers is guarded by the lock
	if s1.sessionTimers.Len() != 0 {
		t.Fatalf("timers should be gone")
	}
}

func TestServer_SessionTTL_Failover(t *testing.T) {
	t.Parallel()
	dir1, s1 := testServer(t)
	defer os.RemoveAll(dir1)
	defer s1.Shutdown()
	testrpc.WaitForTestAgent(t, s1.RPC, "dc1")

	dir2, s2 := testServerDCBootstrap(t, "dc1", false)
	defer os.RemoveAll(dir2)
	defer s2.Shutdown()

	dir3, s3 := testServerDCBootstrap(t, "dc1", false)
	defer os.RemoveAll(dir3)
	defer s3.Shutdown()
	servers := []*Server{s1, s2, s3}

	// Try to join
	joinLAN(t, s2, s1)
	joinLAN(t, s3, s1)
	retry.Run(t, func(r *retry.R) { r.Check(wantPeers(s1, 3)) })

	// Find the leader
	var leader *Server
	for _, s := range servers {
		// Check that s.sessionTimers is empty
		if s.sessionTimers.Len() != 0 {
			t.Fatalf("should have no sessionTimers")
		}
		// Find the leader too
		if s.IsLeader() {
			leader = s
		}
	}
	if leader == nil {
		t.Fatalf("Should have a leader")
	}

	codec := rpcClient(t, leader)
	defer codec.Close()

	// Register a node
	node := structs.RegisterRequest{
		Datacenter: s1.config.Datacenter,
		Node:       "foo",
		Address:    "127.0.0.1",
	}
	var out struct{}
	if err := s1.RPC("Catalog.Register", &node, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a TTL session
	arg := structs.SessionRequest{
		Datacenter: "dc1",
		Op:         structs.SessionCreate,
		Session: structs.Session{
			Node: "foo",
			TTL:  "10s",
		},
	}
	var id1 string
	if err := msgpackrpc.CallWithCodec(codec, "Session.Apply", &arg, &id1); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that sessionTimers has the session ID
	if leader.sessionTimers.Get(id1) == nil {
		t.Fatalf("missing session timer")
	}

	// Shutdown the leader!
	leader.Shutdown()

	// sessionTimers should be cleared on leader shutdown
	if leader.sessionTimers.Len() != 0 {
		t.Fatalf("session timers should be empty on the shutdown leader")
	}
	// Find the new leader
	retry.Run(t, func(r *retry.R) {
		leader = nil
		for _, s := range servers {
			if s.IsLeader() {
				leader = s
			}
		}
		if leader == nil {
			r.Fatal("Should have a new leader")
		}

		// Ensure session timer is restored
		if leader.sessionTimers.Get(id1) == nil {
			r.Fatal("missing session timer")
		}
	})
}
<|endoftext|>


<commit_before>package joingroup

import "github.com/segmentio/kafka-go/protocol"

func init() {
	protocol.Register(&Request{}, &Response{})
}

type Request struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	GroupID            string            `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	SessionTimeoutMS   int32             `kafka:"min=v0,max=v7"`
	RebalanceTimeoutMS int32             `kafka:"min=v1,max=v7"`
	MemberID           string            `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	GroupInstanceID    string            `kafka:"min=v5,max=v5,nullable|min=v6,max=v7,compact,nullable"`
	ProtocolType       string            `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	Protocols          []RequestProtocol `kafka:"min=v0,max=v7"`
}

type RequestProtocol struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	Name     string `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	Metadata []byte `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
}

func (r *Request) ApiKey() protocol.ApiKey {
	return protocol.JoinGroup
}

func (r *Request) Group() string { return r.GroupID }

var _ protocol.GroupMessage = (*Request)(nil)

type Response struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	ThrottleTimeMS int32            `kafka:"min=v2,max=v7"`
	ErrorCode      int16            `kafka:"min=v0,max=v7"`
	GenerationID   int32            `kafka:"min=v0,max=v7"`
	ProtocolName   string           `kafka:"min=v0,max=v5|min=v6,max=v6,compact|min=v7,max=v7,compact,nullable"`
	ProtocolType   string           `kafka:"min=v7,max=v7,compact,nullable"`
	LeaderID       string           `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	MemberID       string           `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	Members        []ResponseMember `kafka:"min=v0,max=v7"`
}

type ResponseMember struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	MemberID        string `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	GroupInstanceID string `kafka:"min=v5,max=v5,nullable|min=v6,max=v7,nullable,compact"`
	Metadata        []byte `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
}

type ResponseMemberMetadata struct{}

func (r *Response) ApiKey() protocol.ApiKey { return protocol.JoinGroup }
<commit_msg>re-order protocol/joingroup response fields (#949)<commit_after>package joingroup

import "github.com/segmentio/kafka-go/protocol"

func init() {
	protocol.Register(&Request{}, &Response{})
}

type Request struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	GroupID            string            `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	SessionTimeoutMS   int32             `kafka:"min=v0,max=v7"`
	RebalanceTimeoutMS int32             `kafka:"min=v1,max=v7"`
	MemberID           string            `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	GroupInstanceID    string            `kafka:"min=v5,max=v5,nullable|min=v6,max=v7,compact,nullable"`
	ProtocolType       string            `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	Protocols          []RequestProtocol `kafka:"min=v0,max=v7"`
}

type RequestProtocol struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	Name     string `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	Metadata []byte `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
}

func (r *Request) ApiKey() protocol.ApiKey {
	return protocol.JoinGroup
}

func (r *Request) Group() string { return r.GroupID }

var _ protocol.GroupMessage = (*Request)(nil)

type Response struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	ThrottleTimeMS int32            `kafka:"min=v2,max=v7"`
	ErrorCode      int16            `kafka:"min=v0,max=v7"`
	GenerationID   int32            `kafka:"min=v0,max=v7"`
	ProtocolType   string           `kafka:"min=v7,max=v7,compact,nullable"`
	ProtocolName   string           `kafka:"min=v0,max=v5|min=v6,max=v6,compact|min=v7,max=v7,compact,nullable"`
	LeaderID       string           `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	MemberID       string           `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	Members        []ResponseMember `kafka:"min=v0,max=v7"`
}

type ResponseMember struct {
	// We need at least one tagged field to indicate that this is a "flexible" message
	// type.
	_ struct{} `kafka:"min=v6,max=v7,tag"`

	MemberID        string `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
	GroupInstanceID string `kafka:"min=v5,max=v5,nullable|min=v6,max=v7,nullable,compact"`
	Metadata        []byte `kafka:"min=v0,max=v5|min=v6,max=v7,compact"`
}

type ResponseMemberMetadata struct{}

func (r *Response) ApiKey() protocol.ApiKey { return protocol.JoinGroup }
<|endoftext|>


<commit_before>package slackboard

const (
	Version = "0.9.0"
)
<commit_msg>bumped version to 0.9.1.<commit_after>package slackboard

const (
	Version = "0.9.1"
)
<|endoftext|>


<commit_before>package tracetcp

import (
	"fmt"
	"net"
	"syscall"
	"time"
)

type implTraceEventType int

const (
	timedOut implTraceEventType = iota
	ttlExpired
	connected
	connectFailed
	errored
)

type implTraceEvent struct {
	evtype    implTraceEventType
	timeStamp time.Time

	localPort  int
	remotePort int
	remoteAddr net.IPAddr
	localAddr  net.IPAddr
	ttl        int
	query      int
	err        error
}

func tryConnect(dest net.IPAddr, port, ttl, query int,
	timeout time.Duration, result chan implTraceEvent) {

	// fill in the event with as much info as we have so far
	event := implTraceEvent{
		remoteAddr: dest,
		remotePort: port,
		ttl:        ttl,
		query:      query,
	}

	returnError := func(err error) {
		event.err = err
		event.evtype = errored
		event.timeStamp = time.Now()
		result <- event
	}

	sock, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
	if err != nil {
		returnError(err)
		return
	}
	defer syscall.Close(sock)

	err = syscall.SetsockoptInt(sock, 0x0, syscall.IP_TTL, ttl)
	if err != nil {
		returnError(err)
		return
	}

	err = syscall.SetNonblock(sock, true)
	if err != nil {
		returnError(err)
		return
	}

	// ignore error from connect in non-blocking mode. as it will always return an
	// in progress error
	_ = syscall.Connect(sock, ToSockaddrInet4(dest, port))

	// get the local ip address and port number
	local, err := syscall.Getsockname(sock)
	if err != nil {
		returnError(err)
		return
	}

	event.localAddr, event.localPort, err = ToIPAddrAndPort(local)
	if err != nil {
		returnError(err)
		return
	}
}

func connect(host string, port, ttl int, timeout time.Duration) error {
	sock, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
	if err != nil {
		return err
	}
	defer syscall.Close(sock)

	err = syscall.SetsockoptInt(sock, 0x0, syscall.IP_TTL, ttl)
	if err != nil {
		return err
	}

	err = syscall.SetNonblock(sock, true)
	if err != nil {
		return err
	}

	addr, err := LookupAddress(host)
	if err != nil {
		return nil
	}

	// ignore error from connect in non-blocking mode. as it will always return a
	// in progress error
	_ = syscall.Connect(sock, &syscall.SockaddrInet4{Port: 80, Addr: addr})

	name, err := syscall.Getsockname(sock)
	fmt.Println(err, name)

	fdset := &syscall.FdSet{}
	timeoutVal := &syscall.Timeval{}
	timeoutVal.Sec = int64(timeout / time.Second)
	timeoutVal.Usec = int64(timeout-time.Duration(timeoutVal.Sec)*time.Second) / 1000

	fmt.Println(timeoutVal)

	FD_ZERO(fdset)
	FD_SET(fdset, sock)

	start := time.Now()
	x, err := syscall.Select(sock+1, nil, fdset, nil, timeoutVal)
	elapsed := time.Since(start)

	fmt.Println(x, elapsed)
	if err != nil {
		return err
	}

	if FD_ISSET(fdset, sock) {
		fmt.Println("conencted")
	} else {
		fmt.Println("timedout")
	}

	return nil
}
<commit_msg>building more connection function & rec ICMP<commit_after>package tracetcp

import (
	//"fmt"
	"net"
	"syscall"
	"time"
)

type implTraceEventType int

const (
	beginConnect implTraceEventType = iota
	timedOut
	ttlExpired
	connected
	connectFailed
	errored
)

type implTraceEvent struct {
	evtype    implTraceEventType
	timeStamp time.Time

	localPort  int
	remotePort int
	remoteAddr net.IPAddr
	localAddr  net.IPAddr
	ttl        int
	query      int
	err        error
}

func makeErrorEvent(event *implTraceEvent, err error) implTraceEvent {
	event.err = err
	event.evtype = errored
	event.timeStamp = time.Now()
	return *event
}

func makeEvent(event *implTraceEvent, evtype implTraceEventType) implTraceEvent {
	event.evtype = evtype
	event.timeStamp = time.Now()
	return *event
}

func tryConnect(dest net.IPAddr, port, ttl, query int,
	timeout time.Duration, result chan implTraceEvent) {

	// fill in the event with as much info as we have so far
	event := implTraceEvent{
		remoteAddr: dest,
		remotePort: port,
		ttl:        ttl,
		query:      query,
	}

	sock, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}
	defer syscall.Close(sock)

	err = syscall.SetsockoptInt(sock, 0x0, syscall.IP_TTL, ttl)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	err = syscall.SetNonblock(sock, true)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	// ignore error from connect in non-blocking mode. as it will always return an
	// in progress error
	_ = syscall.Connect(sock, ToSockaddrInet4(dest, port))

	// get the local ip address and port number
	local, err := syscall.Getsockname(sock)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	// fill in the local endpoint deatils on the event struct
	event.localAddr, event.localPort, err = ToIPAddrAndPort(local)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	result <- makeEvent(&event, beginConnect)

	fdset := &syscall.FdSet{}
	timeoutVal := MakeTimeval(timeout)

	FD_ZERO(fdset)
	FD_SET(fdset, sock)

	_, err = syscall.Select(sock+1, nil, fdset, nil, &timeoutVal)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	// TODO: test for connect failed?

	if FD_ISSET(fdset, sock) {
		// detect if actually connected as select shows ttl expired as connected
		// so if we try to get the remote address and it fails then ttl has expired
		_, err = syscall.Getpeername(sock)
		if err == nil {
			result <- makeEvent(&event, connected)
		} else {
			result <- makeEvent(&event, connectFailed)
		}
	} else {
		result <- makeEvent(&event, timedOut)
	}
}

func receiveICMP(result chan implTraceEvent) {
	event := implTraceEvent{}

	// Set up the socket to receive inbound packets
	sock, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, syscall.IPPROTO_ICMP)
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	err = syscall.Bind(sock, &syscall.SockaddrInet4{})
	if err != nil {
		result <- makeErrorEvent(&event, err)
		return
	}

	var pkt = make([]byte, 1024)
	for {
		_, from, err := syscall.Recvfrom(sock, pkt, 0)
		if err != nil {
			result <- makeErrorEvent(&event, err)
			return
		}

		// fill in the local endpoint deatils on the event struct
		event.localAddr, _, _ = ToIPAddrAndPort(from)
		result <- makeEvent(&event, ttlExpired)
	}
}
<|endoftext|>


<commit_before>package server

import (
	client "github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/deploy/assets"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
)

// Parameters used when creating the kubernetes replication controller in charge
// of a job or pipeline's workers
type workerOptions struct {
	rcName string // Name of the replication controller managing workers

	userImage    string            // The user's pipeline/job image
	labels       map[string]string // k8s labels attached to the Deployment and workers
	parallelism  int32             // Number of replicas the RC maintains
	cacheSize    string            // Size of cache that sidecar uses
	resources    *api.ResourceList // Resources requested by pipeline/job pods
	workerEnv    []api.EnvVar      // Environment vars set in the user container
	volumes      []api.Volume      // Volumes that we expose to the user container
	volumeMounts []api.VolumeMount // Paths where we mount each volume in 'volumes'

	// Secrets that we mount in the worker container (e.g. for reading/writing to
	// s3)
	imagePullSecrets []api.LocalObjectReference
}

func (a *apiServer) workerPodSpec(options *workerOptions) api.PodSpec {
	pullPolicy := a.workerImagePullPolicy
	if pullPolicy == "" {
		pullPolicy = "IfNotPresent"
	}
	// TODO: make the cache sizes configurable
	sidecarEnv := []api.EnvVar{{
		Name:  "BLOCK_CACHE_BYTES",
		Value: options.cacheSize,
	}, {
		Name:  "PFS_CACHE_BYTES",
		Value: "10M",
	}, {
		Name:  "PACH_ROOT",
		Value: a.storageRoot,
	}, {
		Name:  "STORAGE_BACKEND",
		Value: a.storageBackend,
	}}
	// This only happens in local deployment.  We want the workers to be
	// able to read from/write to the hostpath volume as well.
	storageVolumeName := "pach-disk"
	var sidecarVolumeMounts []api.VolumeMount
	if a.storageHostPath != "" {
		options.volumes = append(options.volumes, api.Volume{
			Name: storageVolumeName,
			VolumeSource: api.VolumeSource{
				HostPath: &api.HostPathVolumeSource{
					Path: a.storageHostPath,
				},
			},
		})

		sidecarVolumeMounts = []api.VolumeMount{
			{
				Name:      storageVolumeName,
				MountPath: a.storageRoot,
			},
		}
	}
	secretVolume, secretMount, err := assets.GetSecretVolumeAndMount(a.storageBackend)
	if err == nil {
		options.volumes = append(options.volumes, secretVolume)
		sidecarVolumeMounts = append(sidecarVolumeMounts, secretMount)
	}
	podSpec := api.PodSpec{
		InitContainers: []api.Container{
			{
				Name:            "init",
				Image:           a.workerImage,
				Command:         []string{"/pach/worker.sh"},
				ImagePullPolicy: api.PullPolicy(pullPolicy),
				Env:             options.workerEnv,
				VolumeMounts:    options.volumeMounts,
			},
		},
		Containers: []api.Container{
			{
				Name:    client.PPSWorkerUserContainerName,
				Image:   options.userImage,
				Command: []string{"/pach-bin/guest.sh"},
				SecurityContext: &api.SecurityContext{
					Privileged: &trueVal, // god is this dumb
				},
				ImagePullPolicy: api.PullPolicy(pullPolicy),
				Env:             options.workerEnv,
				VolumeMounts:    options.volumeMounts,
			},
			{
				Name:            client.PPSWorkerSidecarContainerName,
				Image:           a.workerSidecarImage,
				Command:         []string{"/pachd", "--mode", "sidecar"},
				ImagePullPolicy: api.PullPolicy(pullPolicy),
				Env:             sidecarEnv,
				VolumeMounts:    sidecarVolumeMounts,
			},
		},
		RestartPolicy:    "Always",
		Volumes:          options.volumes,
		ImagePullSecrets: options.imagePullSecrets,
	}
	if options.resources != nil {
		podSpec.Containers[0].Resources = api.ResourceRequirements{
			Requests: *options.resources,
		}
	}
	return podSpec
}

func (a *apiServer) getWorkerOptions(rcName string, parallelism int32, resources *api.ResourceList, transform *pps.Transform, cacheSize string) *workerOptions {
	labels := labels(rcName)
	userImage := transform.Image
	if userImage == "" {
		userImage = DefaultUserImage
	}

	var workerEnv []api.EnvVar
	for name, value := range transform.Env {
		workerEnv = append(
			workerEnv,
			api.EnvVar{
				Name:  name,
				Value: value,
			},
		)
	}
	// We use Kubernetes' "Downward API" so the workers know their IP
	// addresses, which they will then post on etcd so the job managers
	// can discover the workers.
	workerEnv = append(workerEnv, api.EnvVar{
		Name: client.PPSWorkerIPEnv,
		ValueFrom: &api.EnvVarSource{
			FieldRef: &api.ObjectFieldSelector{
				APIVersion: "v1",
				FieldPath:  "status.podIP",
			},
		},
	})
	workerEnv = append(workerEnv, api.EnvVar{
		Name: client.PPSPodNameEnv,
		ValueFrom: &api.EnvVarSource{
			FieldRef: &api.ObjectFieldSelector{
				APIVersion: "v1",
				FieldPath:  "metadata.name",
			},
		},
	})
	// Set the etcd prefix env
	workerEnv = append(workerEnv, api.EnvVar{
		Name:  client.PPSEtcdPrefixEnv,
		Value: a.etcdPrefix,
	})
	// Pass along the namespace
	workerEnv = append(workerEnv, api.EnvVar{
		Name:  client.PPSNamespaceEnv,
		Value: a.namespace,
	})

	var volumes []api.Volume
	var volumeMounts []api.VolumeMount
	for _, secret := range transform.Secrets {
		volumes = append(volumes, api.Volume{
			Name: secret.Name,
			VolumeSource: api.VolumeSource{
				Secret: &api.SecretVolumeSource{
					SecretName: secret.Name,
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      secret.Name,
			MountPath: secret.MountPath,
		})
	}

	volumes = append(volumes, api.Volume{
		Name: "pach-bin",
		VolumeSource: api.VolumeSource{
			EmptyDir: &api.EmptyDirVolumeSource{},
		},
	})
	volumeMounts = append(volumeMounts, api.VolumeMount{
		Name:      "pach-bin",
		MountPath: "/pach-bin",
	})

	volumes = append(volumes, api.Volume{
		Name: client.PPSWorkerVolume,
		VolumeSource: api.VolumeSource{
			EmptyDir: &api.EmptyDirVolumeSource{},
		},
	})
	volumeMounts = append(volumeMounts, api.VolumeMount{
		Name:      client.PPSWorkerVolume,
		MountPath: client.PPSScratchSpace,
	})
	if resources != nil && resources.NvidiaGPU() != nil && !resources.NvidiaGPU().IsZero() {
		volumes = append(volumes, api.Volume{
			Name: "root-lib",
			VolumeSource: api.VolumeSource{
				HostPath: &api.HostPathVolumeSource{
					Path: "/usr/lib",
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      "root-lib",
			MountPath: "/rootfs/usr/lib",
		})
	}
	var imagePullSecrets []api.LocalObjectReference
	for _, secret := range transform.ImagePullSecrets {
		imagePullSecrets = append(imagePullSecrets, api.LocalObjectReference{Name: secret})
	}

	return &workerOptions{
		rcName:           rcName,
		labels:           labels,
		parallelism:      int32(parallelism),
		resources:        resources,
		userImage:        userImage,
		workerEnv:        workerEnv,
		volumes:          volumes,
		volumeMounts:     volumeMounts,
		imagePullSecrets: imagePullSecrets,
		cacheSize:        cacheSize,
	}
}

func (a *apiServer) createWorkerRc(options *workerOptions) error {
	rc := &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   options.rcName,
			Labels: options.labels,
		},
		Spec: api.ReplicationControllerSpec{
			Selector: options.labels,
			Replicas: options.parallelism,
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   options.rcName,
					Labels: options.labels,
				},
				Spec: a.workerPodSpec(options),
			},
		},
	}
	if _, err := a.kubeClient.ReplicationControllers(a.namespace).Create(rc); err != nil {
		if !isAlreadyExistsErr(err) {
			return err
		}
	}

	service := &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   options.rcName,
			Labels: options.labels,
		},
		Spec: api.ServiceSpec{
			Selector: options.labels,
			Ports: []api.ServicePort{
				{
					Port: client.PPSWorkerPort,
					Name: "grpc-port",
				},
			},
		},
	}

	if _, err := a.kubeClient.Services(a.namespace).Create(service); err != nil {
		if !isAlreadyExistsErr(err) {
			return err
		}
	}

	return nil
}
<commit_msg>Mount object store credentials in the user container so that the user container can talk to the object store.<commit_after>package server

import (
	client "github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/deploy/assets"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
)

// Parameters used when creating the kubernetes replication controller in charge
// of a job or pipeline's workers
type workerOptions struct {
	rcName string // Name of the replication controller managing workers

	userImage    string            // The user's pipeline/job image
	labels       map[string]string // k8s labels attached to the Deployment and workers
	parallelism  int32             // Number of replicas the RC maintains
	cacheSize    string            // Size of cache that sidecar uses
	resources    *api.ResourceList // Resources requested by pipeline/job pods
	workerEnv    []api.EnvVar      // Environment vars set in the user container
	volumes      []api.Volume      // Volumes that we expose to the user container
	volumeMounts []api.VolumeMount // Paths where we mount each volume in 'volumes'

	// Secrets that we mount in the worker container (e.g. for reading/writing to
	// s3)
	imagePullSecrets []api.LocalObjectReference
}

func (a *apiServer) workerPodSpec(options *workerOptions) api.PodSpec {
	pullPolicy := a.workerImagePullPolicy
	if pullPolicy == "" {
		pullPolicy = "IfNotPresent"
	}
	// TODO: make the cache sizes configurable
	sidecarEnv := []api.EnvVar{{
		Name:  "BLOCK_CACHE_BYTES",
		Value: options.cacheSize,
	}, {
		Name:  "PFS_CACHE_BYTES",
		Value: "10M",
	}, {
		Name:  "PACH_ROOT",
		Value: a.storageRoot,
	}, {
		Name:  "STORAGE_BACKEND",
		Value: a.storageBackend,
	}}
	// This only happens in local deployment.  We want the workers to be
	// able to read from/write to the hostpath volume as well.
	storageVolumeName := "pach-disk"
	var sidecarVolumeMounts []api.VolumeMount
	if a.storageHostPath != "" {
		options.volumes = append(options.volumes, api.Volume{
			Name: storageVolumeName,
			VolumeSource: api.VolumeSource{
				HostPath: &api.HostPathVolumeSource{
					Path: a.storageHostPath,
				},
			},
		})

		sidecarVolumeMounts = []api.VolumeMount{
			{
				Name:      storageVolumeName,
				MountPath: a.storageRoot,
			},
		}
	}
	secretVolume, secretMount, err := assets.GetSecretVolumeAndMount(a.storageBackend)
	if err == nil {
		options.volumes = append(options.volumes, secretVolume)
		options.volumeMounts = append(options.volumeMounts, secretMount)
		sidecarVolumeMounts = append(sidecarVolumeMounts, secretMount)
	}
	podSpec := api.PodSpec{
		InitContainers: []api.Container{
			{
				Name:            "init",
				Image:           a.workerImage,
				Command:         []string{"/pach/worker.sh"},
				ImagePullPolicy: api.PullPolicy(pullPolicy),
			},
		},
		Containers: []api.Container{
			{
				Name:    client.PPSWorkerUserContainerName,
				Image:   options.userImage,
				Command: []string{"/pach-bin/guest.sh"},
				SecurityContext: &api.SecurityContext{
					Privileged: &trueVal, // god is this dumb
				},
				ImagePullPolicy: api.PullPolicy(pullPolicy),
				Env:             options.workerEnv,
				VolumeMounts:    options.volumeMounts,
			},
			{
				Name:            client.PPSWorkerSidecarContainerName,
				Image:           a.workerSidecarImage,
				Command:         []string{"/pachd", "--mode", "sidecar"},
				ImagePullPolicy: api.PullPolicy(pullPolicy),
				Env:             sidecarEnv,
				VolumeMounts:    sidecarVolumeMounts,
			},
		},
		RestartPolicy:    "Always",
		Volumes:          options.volumes,
		ImagePullSecrets: options.imagePullSecrets,
	}
	if options.resources != nil {
		podSpec.Containers[0].Resources = api.ResourceRequirements{
			Requests: *options.resources,
		}
	}
	return podSpec
}

func (a *apiServer) getWorkerOptions(rcName string, parallelism int32, resources *api.ResourceList, transform *pps.Transform, cacheSize string) *workerOptions {
	labels := labels(rcName)
	userImage := transform.Image
	if userImage == "" {
		userImage = DefaultUserImage
	}

	var workerEnv []api.EnvVar
	for name, value := range transform.Env {
		workerEnv = append(
			workerEnv,
			api.EnvVar{
				Name:  name,
				Value: value,
			},
		)
	}
	// We use Kubernetes' "Downward API" so the workers know their IP
	// addresses, which they will then post on etcd so the job managers
	// can discover the workers.
	workerEnv = append(workerEnv, api.EnvVar{
		Name: client.PPSWorkerIPEnv,
		ValueFrom: &api.EnvVarSource{
			FieldRef: &api.ObjectFieldSelector{
				APIVersion: "v1",
				FieldPath:  "status.podIP",
			},
		},
	})
	workerEnv = append(workerEnv, api.EnvVar{
		Name: client.PPSPodNameEnv,
		ValueFrom: &api.EnvVarSource{
			FieldRef: &api.ObjectFieldSelector{
				APIVersion: "v1",
				FieldPath:  "metadata.name",
			},
		},
	})
	// Set the etcd prefix env
	workerEnv = append(workerEnv, api.EnvVar{
		Name:  client.PPSEtcdPrefixEnv,
		Value: a.etcdPrefix,
	})
	// Pass along the namespace
	workerEnv = append(workerEnv, api.EnvVar{
		Name:  client.PPSNamespaceEnv,
		Value: a.namespace,
	})

	var volumes []api.Volume
	var volumeMounts []api.VolumeMount
	for _, secret := range transform.Secrets {
		volumes = append(volumes, api.Volume{
			Name: secret.Name,
			VolumeSource: api.VolumeSource{
				Secret: &api.SecretVolumeSource{
					SecretName: secret.Name,
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      secret.Name,
			MountPath: secret.MountPath,
		})
	}

	volumes = append(volumes, api.Volume{
		Name: "pach-bin",
		VolumeSource: api.VolumeSource{
			EmptyDir: &api.EmptyDirVolumeSource{},
		},
	})
	volumeMounts = append(volumeMounts, api.VolumeMount{
		Name:      "pach-bin",
		MountPath: "/pach-bin",
	})

	volumes = append(volumes, api.Volume{
		Name: client.PPSWorkerVolume,
		VolumeSource: api.VolumeSource{
			EmptyDir: &api.EmptyDirVolumeSource{},
		},
	})
	volumeMounts = append(volumeMounts, api.VolumeMount{
		Name:      client.PPSWorkerVolume,
		MountPath: client.PPSScratchSpace,
	})
	if resources != nil && resources.NvidiaGPU() != nil && !resources.NvidiaGPU().IsZero() {
		volumes = append(volumes, api.Volume{
			Name: "root-lib",
			VolumeSource: api.VolumeSource{
				HostPath: &api.HostPathVolumeSource{
					Path: "/usr/lib",
				},
			},
		})
		volumeMounts = append(volumeMounts, api.VolumeMount{
			Name:      "root-lib",
			MountPath: "/rootfs/usr/lib",
		})
	}
	var imagePullSecrets []api.LocalObjectReference
	for _, secret := range transform.ImagePullSecrets {
		imagePullSecrets = append(imagePullSecrets, api.LocalObjectReference{Name: secret})
	}

	return &workerOptions{
		rcName:           rcName,
		labels:           labels,
		parallelism:      int32(parallelism),
		resources:        resources,
		userImage:        userImage,
		workerEnv:        workerEnv,
		volumes:          volumes,
		volumeMounts:     volumeMounts,
		imagePullSecrets: imagePullSecrets,
		cacheSize:        cacheSize,
	}
}

func (a *apiServer) createWorkerRc(options *workerOptions) error {
	rc := &api.ReplicationController{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "ReplicationController",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   options.rcName,
			Labels: options.labels,
		},
		Spec: api.ReplicationControllerSpec{
			Selector: options.labels,
			Replicas: options.parallelism,
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   options.rcName,
					Labels: options.labels,
				},
				Spec: a.workerPodSpec(options),
			},
		},
	}
	if _, err := a.kubeClient.ReplicationControllers(a.namespace).Create(rc); err != nil {
		if !isAlreadyExistsErr(err) {
			return err
		}
	}

	service := &api.Service{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   options.rcName,
			Labels: options.labels,
		},
		Spec: api.ServiceSpec{
			Selector: options.labels,
			Ports: []api.ServicePort{
				{
					Port: client.PPSWorkerPort,
					Name: "grpc-port",
				},
			},
		},
	}

	if _, err := a.kubeClient.Services(a.namespace).Create(service); err != nil {
		if !isAlreadyExistsErr(err) {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package resourcequota

import (
	"encoding/json"
	"reflect"
	"sync"
	"time"

	namespaceutil "github.com/rancher/rancher/pkg/namespace"
	validate "github.com/rancher/rancher/pkg/resourcequota"
	"github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/cache"
	clientcache "k8s.io/client-go/tools/cache"
)

const (
	projectIDAnnotation             = "field.cattle.io/projectId"
	resourceQuotaLabel              = "resourcequota.management.cattle.io/default-resource-quota"
	resourceQuotaInitCondition      = "ResourceQuotaInit"
	resourceQuotaAnnotation         = "field.cattle.io/resourceQuota"
	resourceQuotaValidatedCondition = "ResourceQuotaValidated"
)

var (
	projectLockCache = cache.NewLRUExpireCache(1000)
)

/*
SyncController takes care of creating Kubernetes resource quota based on the resource limits
defined in namespace.resourceQuota
*/
type SyncController struct {
	ProjectLister       v3.ProjectLister
	Namespaces          v1.NamespaceInterface
	NamespaceLister     v1.NamespaceLister
	ResourceQuotas      v1.ResourceQuotaInterface
	ResourceQuotaLister v1.ResourceQuotaLister
	NsIndexer           clientcache.Indexer
}

func (c *SyncController) syncResourceQuota(key string, ns *corev1.Namespace) error {
	if ns == nil || ns.DeletionTimestamp != nil {
		return nil
	}

	_, err := c.CreateResourceQuota(ns)
	return err
}

func (c *SyncController) CreateResourceQuota(ns *corev1.Namespace) (*corev1.Namespace, error) {
	existing, err := c.getExistingResourceQuota(ns)
	if err != nil {
		return ns, err
	}

	projectLimit, _, err := getProjectResourceQuotaLimit(ns, c.ProjectLister)
	if err != nil {
		return ns, err
	}

	var quotaSpec *corev1.ResourceQuotaSpec
	if projectLimit != nil {
		quotaSpec, err = c.getNamespaceResourceQuota(ns, true)
		if err != nil {
			return ns, err
		}
	}

	operation := "none"
	if existing == nil {
		if quotaSpec != nil {
			operation = "create"
		}
	} else {
		if quotaSpec == nil {
			operation = "delete"
		} else if !reflect.DeepEqual(existing.Spec.Hard, quotaSpec.Hard) {
			operation = "update"
		}
	}

	var updated *corev1.Namespace
	var isFit bool
	switch operation {
	case "create":
		isFit, updated, err = c.validateAndSetNamespaceQuota(ns)
		if err != nil || !isFit {
			return updated, err
		}
		err = c.createDefaultResourceQuota(ns, quotaSpec)
	case "update":
		isFit, updated, err = c.validateAndSetNamespaceQuota(ns)
		if err != nil || !isFit {
			return updated, err
		}
		err = c.updateResourceQuota(existing, quotaSpec)
	case "delete":
		err = c.deleteResourceQuota(existing)
	}

	toReturn := updated
	if toReturn == nil {
		toReturn = ns
	}

	if err == nil {
		set, err := namespaceutil.IsNamespaceConditionSet(ns, resourceQuotaInitCondition, true)
		if err != nil || set {
			return toReturn, err
		}
		namespaceutil.SetNamespaceCondition(toReturn, time.Second*1, resourceQuotaInitCondition, true, "")
		return c.Namespaces.Update(toReturn)
	}

	return toReturn, err
}

func (c *SyncController) updateResourceQuota(quota *corev1.ResourceQuota, spec *corev1.ResourceQuotaSpec) error {
	toUpdate := quota.DeepCopy()
	toUpdate.Spec = *spec
	logrus.Infof("Updating default resource quota for namespace %v", toUpdate.Namespace)
	_, err := c.ResourceQuotas.Update(toUpdate)
	return err
}

func (c *SyncController) deleteResourceQuota(quota *corev1.ResourceQuota) error {
	logrus.Infof("Deleting default resource quota for namespace %v", quota.Namespace)
	return c.ResourceQuotas.DeleteNamespaced(quota.Namespace, quota.Name, &metav1.DeleteOptions{})
}

func (c *SyncController) getExistingResourceQuota(ns *corev1.Namespace) (*corev1.ResourceQuota, error) {
	set := labels.Set(map[string]string{resourceQuotaLabel: "true"})
	quota, err := c.ResourceQuotaLister.List(ns.Name, set.AsSelector())
	if err != nil {
		return nil, err
	}
	if len(quota) == 0 {
		return nil, nil
	}
	return quota[0], nil
}

func (c *SyncController) getNamespaceResourceQuota(ns *corev1.Namespace, setDefault bool) (*corev1.ResourceQuotaSpec, error) {
	limit, err := getNamespaceLimit(ns)
	if err != nil {
		return nil, err
	}
	if limit == nil {
		if setDefault {
			limit = defaultResourceLimit
		} else {
			return nil, nil
		}
	}

	return convertResourceLimitResourceQuotaSpec(limit)
}

var defaultResourceLimit = &v3.ResourceQuotaLimit{
	Pods:                   "0",
	Services:               "0",
	ReplicationControllers: "0",
	Secrets:                "0",
	ConfigMaps:             "0",
	PersistentVolumeClaims: "0",
	ServicesNodePorts:      "0",
	ServicesLoadBalancers:  "0",
	RequestsCPU:            "0",
	RequestsMemory:         "0",
	RequestsStorage:        "0",
	LimitsCPU:              "0",
	LimitsMemory:           "0",
}

func (c *SyncController) createDefaultResourceQuota(ns *corev1.Namespace, spec *corev1.ResourceQuotaSpec) error {
	resourceQuota := &corev1.ResourceQuota{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "default-",
			Namespace:    ns.Name,
			Labels:       map[string]string{resourceQuotaLabel: "true"},
		},
		Spec: *spec,
	}
	logrus.Infof("Creating default resource quota for namespace %v", ns.Name)
	_, err := c.ResourceQuotas.Create(resourceQuota)
	return err
}

func (c *SyncController) validateAndSetNamespaceQuota(ns *corev1.Namespace) (bool, *corev1.Namespace, error) {
	if ns == nil || ns.DeletionTimestamp != nil {
		return true, ns, nil
	}

	// get project limit
	projectLimit, projectID, err := getProjectResourceQuotaLimit(ns, c.ProjectLister)
	if err != nil {
		return false, ns, err
	}

	if projectLimit == nil {
		return true, ns, err
	}

	// set default quota if not set
	quotaToUpdate, err := c.getResourceQuotaToUpdate(ns)
	if err != nil {
		return false, ns, err
	}
	updatedNs := ns.DeepCopy()
	if quotaToUpdate != "" {
		if updatedNs.Annotations == nil {
			updatedNs.Annotations = map[string]string{}
		}
		updatedNs.Annotations[resourceQuotaAnnotation] = quotaToUpdate
		updatedNs, err = c.Namespaces.Update(updatedNs)
		if err != nil {
			return false, updatedNs, err
		}
	}

	// validate resource quota
	mu := getProjectLock(projectID)
	mu.Lock()
	defer mu.Unlock()
	// get other Namespaces
	objects, err := c.NsIndexer.ByIndex(nsByProjectIndex, projectID)
	if err != nil {
		return false, updatedNs, err
	}
	var nsLimits []*v3.ResourceQuotaLimit
	for _, o := range objects {
		n := o.(*corev1.Namespace)
		// skip itself
		if n.Name == ns.Name {
			continue
		}
		nsLimit, err := getNamespaceLimit(ns)
		if err != nil {
			return false, updatedNs, err
		}
		nsLimits = append(nsLimits, nsLimit)
	}
	nsLimit, err := getNamespaceLimit(updatedNs)
	if err != nil {
		return false, updatedNs, err
	}
	isFit, msg, err := validate.IsQuotaFit(nsLimit, nsLimits, projectLimit)
	if err != nil {
		return false, updatedNs, err
	}

	validated, err := c.setValidated(updatedNs, isFit, msg)

	return isFit, validated, err

}

func (c *SyncController) setValidated(ns *corev1.Namespace, value bool, msg string) (*corev1.Namespace, error) {
	set, err := namespaceutil.IsNamespaceConditionSet(ns, resourceQuotaValidatedCondition, value)
	if set || err != nil {
		return ns, err
	}
	toUpdate := ns.DeepCopy()
	err = namespaceutil.SetNamespaceCondition(toUpdate, time.Second*1, resourceQuotaValidatedCondition, value, msg)
	if err != nil {
		return ns, err
	}
	return c.Namespaces.Update(toUpdate)
}

func getProjectLock(projectID string) *sync.Mutex {
	val, ok := projectLockCache.Get(projectID)
	if !ok {
		projectLockCache.Add(projectID, &sync.Mutex{}, time.Hour)
		val, _ = projectLockCache.Get(projectID)
	}
	mu := val.(*sync.Mutex)
	return mu
}

func (c *SyncController) getResourceQuotaToUpdate(ns *corev1.Namespace) (string, error) {
	if getNamespaceResourceQuota(ns) != "" {
		return "", nil
	}

	quota, err := getProjectNamespaceDefaultQuota(ns, c.ProjectLister)
	if err != nil {
		return "", err
	}
	b, err := json.Marshal(quota)
	if err != nil {
		return "", err
	}
	return string(b), nil

}
<commit_msg>Create default resource quota if user passed quota is not a fit<commit_after>package resourcequota

import (
	"encoding/json"
	"reflect"
	"sync"
	"time"

	namespaceutil "github.com/rancher/rancher/pkg/namespace"
	validate "github.com/rancher/rancher/pkg/resourcequota"
	"github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/cache"
	clientcache "k8s.io/client-go/tools/cache"
)

const (
	projectIDAnnotation             = "field.cattle.io/projectId"
	resourceQuotaLabel              = "resourcequota.management.cattle.io/default-resource-quota"
	resourceQuotaInitCondition      = "ResourceQuotaInit"
	resourceQuotaAnnotation         = "field.cattle.io/resourceQuota"
	resourceQuotaValidatedCondition = "ResourceQuotaValidated"
)

var (
	projectLockCache = cache.NewLRUExpireCache(1000)
)

/*
SyncController takes care of creating Kubernetes resource quota based on the resource limits
defined in namespace.resourceQuota
*/
type SyncController struct {
	ProjectLister       v3.ProjectLister
	Namespaces          v1.NamespaceInterface
	NamespaceLister     v1.NamespaceLister
	ResourceQuotas      v1.ResourceQuotaInterface
	ResourceQuotaLister v1.ResourceQuotaLister
	NsIndexer           clientcache.Indexer
}

func (c *SyncController) syncResourceQuota(key string, ns *corev1.Namespace) error {
	if ns == nil || ns.DeletionTimestamp != nil {
		return nil
	}

	_, err := c.CreateResourceQuota(ns)
	return err
}

func (c *SyncController) CreateResourceQuota(ns *corev1.Namespace) (*corev1.Namespace, error) {
	existing, err := c.getExistingResourceQuota(ns)
	if err != nil {
		return ns, err
	}

	projectLimit, _, err := getProjectResourceQuotaLimit(ns, c.ProjectLister)
	if err != nil {
		return ns, err
	}

	var quotaSpec *corev1.ResourceQuotaSpec
	if projectLimit != nil {
		quotaSpec, err = c.getNamespaceResourceQuota(ns, true)
		if err != nil {
			return ns, err
		}
	}

	operation := "none"
	if existing == nil {
		if quotaSpec != nil {
			operation = "create"
		}
	} else {
		if quotaSpec == nil {
			operation = "delete"
		} else if !reflect.DeepEqual(existing.Spec.Hard, quotaSpec.Hard) {
			operation = "update"
		}
	}

	var updated *corev1.Namespace
	var isFit bool
	switch operation {
	case "create":
		isFit, updated, err = c.validateAndSetNamespaceQuota(ns)
		if err != nil {
			return updated, err
		}
		if !isFit {
			// create default "all 0" resource quota
			quotaSpec, err = getDefaultQuotaSpec()
			if err != nil {
				return updated, err
			}
		}
		err = c.createDefaultResourceQuota(ns, quotaSpec)
	case "update":
		isFit, updated, err = c.validateAndSetNamespaceQuota(ns)
		if err != nil || !isFit {
			return updated, err
		}
		err = c.updateResourceQuota(existing, quotaSpec)
	case "delete":
		err = c.deleteResourceQuota(existing)
	}

	toReturn := updated
	if toReturn == nil {
		toReturn = ns
	}

	if err == nil {
		set, err := namespaceutil.IsNamespaceConditionSet(ns, resourceQuotaInitCondition, true)
		if err != nil || set {
			return toReturn, err
		}
		namespaceutil.SetNamespaceCondition(toReturn, time.Second*1, resourceQuotaInitCondition, true, "")
		return c.Namespaces.Update(toReturn)
	}

	return toReturn, err
}

func (c *SyncController) updateResourceQuota(quota *corev1.ResourceQuota, spec *corev1.ResourceQuotaSpec) error {
	toUpdate := quota.DeepCopy()
	toUpdate.Spec = *spec
	logrus.Infof("Updating default resource quota for namespace %v", toUpdate.Namespace)
	_, err := c.ResourceQuotas.Update(toUpdate)
	return err
}

func (c *SyncController) deleteResourceQuota(quota *corev1.ResourceQuota) error {
	logrus.Infof("Deleting default resource quota for namespace %v", quota.Namespace)
	return c.ResourceQuotas.DeleteNamespaced(quota.Namespace, quota.Name, &metav1.DeleteOptions{})
}

func (c *SyncController) getExistingResourceQuota(ns *corev1.Namespace) (*corev1.ResourceQuota, error) {
	set := labels.Set(map[string]string{resourceQuotaLabel: "true"})
	quota, err := c.ResourceQuotaLister.List(ns.Name, set.AsSelector())
	if err != nil {
		return nil, err
	}
	if len(quota) == 0 {
		return nil, nil
	}
	return quota[0], nil
}

func (c *SyncController) getNamespaceResourceQuota(ns *corev1.Namespace, setDefault bool) (*corev1.ResourceQuotaSpec, error) {
	limit, err := getNamespaceLimit(ns)
	if err != nil {
		return nil, err
	}
	if limit == nil {
		if setDefault {
			limit = defaultResourceLimit
		} else {
			return nil, nil
		}
	}

	return convertResourceLimitResourceQuotaSpec(limit)
}

func getDefaultQuotaSpec() (*corev1.ResourceQuotaSpec, error) {
	return convertResourceLimitResourceQuotaSpec(defaultResourceLimit)
}

var defaultResourceLimit = &v3.ResourceQuotaLimit{
	Pods:                   "0",
	Services:               "0",
	ReplicationControllers: "0",
	Secrets:                "0",
	ConfigMaps:             "0",
	PersistentVolumeClaims: "0",
	ServicesNodePorts:      "0",
	ServicesLoadBalancers:  "0",
	RequestsCPU:            "0",
	RequestsMemory:         "0",
	RequestsStorage:        "0",
	LimitsCPU:              "0",
	LimitsMemory:           "0",
}

func (c *SyncController) createDefaultResourceQuota(ns *corev1.Namespace, spec *corev1.ResourceQuotaSpec) error {
	resourceQuota := &corev1.ResourceQuota{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "default-",
			Namespace:    ns.Name,
			Labels:       map[string]string{resourceQuotaLabel: "true"},
		},
		Spec: *spec,
	}
	logrus.Infof("Creating default resource quota for namespace %v", ns.Name)
	_, err := c.ResourceQuotas.Create(resourceQuota)
	return err
}

func (c *SyncController) validateAndSetNamespaceQuota(ns *corev1.Namespace) (bool, *corev1.Namespace, error) {
	if ns == nil || ns.DeletionTimestamp != nil {
		return true, ns, nil
	}

	// get project limit
	projectLimit, projectID, err := getProjectResourceQuotaLimit(ns, c.ProjectLister)
	if err != nil {
		return false, ns, err
	}

	if projectLimit == nil {
		return true, ns, err
	}

	// set default quota if not set
	quotaToUpdate, err := c.getResourceQuotaToUpdate(ns)
	if err != nil {
		return false, ns, err
	}
	updatedNs := ns.DeepCopy()
	if quotaToUpdate != "" {
		if updatedNs.Annotations == nil {
			updatedNs.Annotations = map[string]string{}
		}
		updatedNs.Annotations[resourceQuotaAnnotation] = quotaToUpdate
		updatedNs, err = c.Namespaces.Update(updatedNs)
		if err != nil {
			return false, updatedNs, err
		}
	}

	// validate resource quota
	mu := getProjectLock(projectID)
	mu.Lock()
	defer mu.Unlock()
	// get other Namespaces
	objects, err := c.NsIndexer.ByIndex(nsByProjectIndex, projectID)
	if err != nil {
		return false, updatedNs, err
	}
	var nsLimits []*v3.ResourceQuotaLimit
	for _, o := range objects {
		n := o.(*corev1.Namespace)
		// skip itself
		if n.Name == ns.Name {
			continue
		}
		nsLimit, err := getNamespaceLimit(ns)
		if err != nil {
			return false, updatedNs, err
		}
		nsLimits = append(nsLimits, nsLimit)
	}
	nsLimit, err := getNamespaceLimit(updatedNs)
	if err != nil {
		return false, updatedNs, err
	}
	isFit, msg, err := validate.IsQuotaFit(nsLimit, nsLimits, projectLimit)
	if err != nil {
		return false, updatedNs, err
	}

	validated, err := c.setValidated(updatedNs, isFit, msg)

	return isFit, validated, err

}

func (c *SyncController) setValidated(ns *corev1.Namespace, value bool, msg string) (*corev1.Namespace, error) {
	set, err := namespaceutil.IsNamespaceConditionSet(ns, resourceQuotaValidatedCondition, value)
	if set || err != nil {
		return ns, err
	}
	toUpdate := ns.DeepCopy()
	err = namespaceutil.SetNamespaceCondition(toUpdate, time.Second*1, resourceQuotaValidatedCondition, value, msg)
	if err != nil {
		return ns, err
	}
	return c.Namespaces.Update(toUpdate)
}

func getProjectLock(projectID string) *sync.Mutex {
	val, ok := projectLockCache.Get(projectID)
	if !ok {
		projectLockCache.Add(projectID, &sync.Mutex{}, time.Hour)
		val, _ = projectLockCache.Get(projectID)
	}
	mu := val.(*sync.Mutex)
	return mu
}

func (c *SyncController) getResourceQuotaToUpdate(ns *corev1.Namespace) (string, error) {
	if getNamespaceResourceQuota(ns) != "" {
		return "", nil
	}

	quota, err := getProjectNamespaceDefaultQuota(ns, c.ProjectLister)
	if err != nil {
		return "", err
	}
	b, err := json.Marshal(quota)
	if err != nil {
		return "", err
	}
	return string(b), nil

}
<|endoftext|>


<commit_before>package oauth

import (
	"fmt"
	"os"
)

// Stores access token information to the given filename.
//
// Format is the same as the return from the server.
func (o *OAuth) Save(fileName string) (err error) {
	file, err := os.OpenFile(fileName, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0600)
	defer file.Close()
	if err != nil {
		return
	}

	fmt.Fprintf(file, "oauth_token=%s&oauth_token_secret=%s", o.accessToken, o.accessSecret)
	if o.userId != 0 {
		fmt.Fprintf(file, "&user_id=%d", o.userId)
	}
	if o.userName != "" {
		fmt.Fprintf(file, "&screen_name=%s", o.userName)
	}

	return nil
}

// Loads access token information from a file.
func (o *OAuth) Load(fileName string) (err error) {
	file, err := os.Open(fileName)
	defer file.Close()
	if err != nil {
		return
	}

	return o.parseResponse(200, file, TokenReq)
}
<commit_msg>Fix file close defers to not try to close nil.<commit_after>package oauth

import (
	"fmt"
	"os"
)

// Stores access token information to the given filename.
//
// Format is the same as the return from the server.
func (o *OAuth) Save(fileName string) (err error) {
	file, err := os.OpenFile(fileName, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0600)
	if err != nil {
		return
	}
	defer file.Close()

	fmt.Fprintf(file, "oauth_token=%s&oauth_token_secret=%s", o.accessToken, o.accessSecret)
	if o.userId != 0 {
		fmt.Fprintf(file, "&user_id=%d", o.userId)
	}
	if o.userName != "" {
		fmt.Fprintf(file, "&screen_name=%s", o.userName)
	}

	return nil
}

// Loads access token information from a file.
func (o *OAuth) Load(fileName string) (err error) {
	file, err := os.Open(fileName)
	if err != nil {
		return
	}
	defer file.Close()

	return o.parseResponse(200, file, TokenReq)
}
<|endoftext|>


<commit_before>package birc

import (
	"crypto/tls"
	"fmt"
	"hash/crc32"
	"net"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/42wim/matterbridge/bridge"
	"github.com/42wim/matterbridge/bridge/config"
	"github.com/42wim/matterbridge/bridge/helper"
	"github.com/lrstanley/girc"

	// We need to import the 'data' package as an implicit dependency.
	// See: https://godoc.org/github.com/paulrosania/go-charset/charset
	_ "github.com/paulrosania/go-charset/data"
)

type Birc struct {
	i                                         *girc.Client
	Nick                                      string
	names                                     map[string][]string
	connected                                 chan error
	Local                                     chan config.Message // local queue for flood control
	FirstConnection, authDone                 bool
	MessageDelay, MessageQueue, MessageLength int

	*bridge.Config
}

func New(cfg *bridge.Config) bridge.Bridger {
	b := &Birc{}
	b.Config = cfg
	b.Nick = b.GetString("Nick")
	b.names = make(map[string][]string)
	b.connected = make(chan error)
	if b.GetInt("MessageDelay") == 0 {
		b.MessageDelay = 1300
	} else {
		b.MessageDelay = b.GetInt("MessageDelay")
	}
	if b.GetInt("MessageQueue") == 0 {
		b.MessageQueue = 30
	} else {
		b.MessageQueue = b.GetInt("MessageQueue")
	}
	if b.GetInt("MessageLength") == 0 {
		b.MessageLength = 400
	} else {
		b.MessageLength = b.GetInt("MessageLength")
	}
	b.FirstConnection = true
	return b
}

func (b *Birc) Command(msg *config.Message) string {
	if msg.Text == "!users" {
		b.i.Handlers.Add(girc.RPL_NAMREPLY, b.storeNames)
		b.i.Handlers.Add(girc.RPL_ENDOFNAMES, b.endNames)
		b.i.Cmd.SendRaw("NAMES " + msg.Channel) //nolint:errcheck
	}
	return ""
}

func (b *Birc) Connect() error {
	b.Local = make(chan config.Message, b.MessageQueue+10)
	b.Log.Infof("Connecting %s", b.GetString("Server"))

	i, err := b.getClient()
	if err != nil {
		return err
	}

	if b.GetBool("UseSASL") {
		i.Config.SASL = &girc.SASLPlain{
			User: b.GetString("NickServNick"),
			Pass: b.GetString("NickServPassword"),
		}
	}

	i.Handlers.Add(girc.RPL_WELCOME, b.handleNewConnection)
	i.Handlers.Add(girc.RPL_ENDOFMOTD, b.handleOtherAuth)
	i.Handlers.Add(girc.ERR_NOMOTD, b.handleOtherAuth)
	i.Handlers.Add(girc.ALL_EVENTS, b.handleOther)
	b.i = i

	go b.doConnect()

	err = <-b.connected
	if err != nil {
		return fmt.Errorf("connection failed %s", err)
	}
	b.Log.Info("Connection succeeded")
	b.FirstConnection = false
	if b.GetInt("DebugLevel") == 0 {
		i.Handlers.Clear(girc.ALL_EVENTS)
	}
	go b.doSend()
	return nil
}

func (b *Birc) Disconnect() error {
	b.i.Close()
	close(b.Local)
	return nil
}

func (b *Birc) JoinChannel(channel config.ChannelInfo) error {
	// need to check if we have nickserv auth done before joining channels
	for {
		if b.authDone {
			break
		}
		time.Sleep(time.Second)
	}
	if channel.Options.Key != "" {
		b.Log.Debugf("using key %s for channel %s", channel.Options.Key, channel.Name)
		b.i.Cmd.JoinKey(channel.Name, channel.Options.Key)
	} else {
		b.i.Cmd.Join(channel.Name)
	}
	return nil
}

func (b *Birc) Send(msg config.Message) (string, error) {
	// ignore delete messages
	if msg.Event == config.EventMsgDelete {
		return "", nil
	}

	b.Log.Debugf("=> Receiving %#v", msg)

	// we can be in between reconnects #385
	if !b.i.IsConnected() {
		b.Log.Error("Not connected to server, dropping message")
		return "", nil
	}

	// Execute a command
	if strings.HasPrefix(msg.Text, "!") {
		b.Command(&msg)
	}

	// convert to specified charset
	if err := b.handleCharset(&msg); err != nil {
		return "", err
	}

	// handle files, return if we're done here
	if ok := b.handleFiles(&msg); ok {
		return "", nil
	}

	var msgLines []string
	if b.GetBool("MessageSplit") {
		msgLines = helper.GetSubLines(msg.Text, b.MessageLength)
	} else {
		msgLines = helper.GetSubLines(msg.Text, 0)
	}
	for i := range msgLines {
		if len(b.Local) >= b.MessageQueue {
			b.Log.Debugf("flooding, dropping message (queue at %d)", len(b.Local))
			return "", nil
		}

		b.Local <- config.Message{
			Text:     msgLines[i],
			Username: msg.Username,
			Channel:  msg.Channel,
			Event:    msg.Event,
		}
	}
	return "", nil
}

func (b *Birc) doConnect() {
	for {
		if err := b.i.Connect(); err != nil {
			b.Log.Errorf("disconnect: error: %s", err)
			if b.FirstConnection {
				b.connected <- err
				return
			}
		} else {
			b.Log.Info("disconnect: client requested quit")
		}
		b.Log.Info("reconnecting in 30 seconds...")
		time.Sleep(30 * time.Second)
		b.i.Handlers.Clear(girc.RPL_WELCOME)
		b.i.Handlers.Add(girc.RPL_WELCOME, func(client *girc.Client, event girc.Event) {
			b.Remote <- config.Message{Username: "system", Text: "rejoin", Channel: "", Account: b.Account, Event: config.EventRejoinChannels}
			// set our correct nick on reconnect if necessary
			b.Nick = event.Source.Name
		})
	}
}

func (b *Birc) doSend() {
	rate := time.Millisecond * time.Duration(b.MessageDelay)
	throttle := time.NewTicker(rate)
	for msg := range b.Local {
		<-throttle.C
		username := msg.Username
		if b.GetBool("Colornicks") {
			checksum := crc32.ChecksumIEEE([]byte(msg.Username))
			colorCode := checksum%14 + 2 // quick fix - prevent white or black color codes
			username = fmt.Sprintf("\x03%02d%s\x0F", colorCode, msg.Username)
		}
		if msg.Event == config.EventUserAction {
			b.i.Cmd.Action(msg.Channel, username+msg.Text)
		} else {
			b.Log.Debugf("Sending to channel %s", msg.Channel)
			b.i.Cmd.Message(msg.Channel, username+msg.Text)
		}
	}
}

// validateInput validates the server/port/nick configuration. Returns a *girc.Client if successful
func (b *Birc) getClient() (*girc.Client, error) {
	server, portstr, err := net.SplitHostPort(b.GetString("Server"))
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portstr)
	if err != nil {
		return nil, err
	}
	// fix strict user handling of girc
	user := b.GetString("Nick")
	for !girc.IsValidUser(user) {
		if len(user) == 1 || len(user) == 0 {
			user = "matterbridge"
			break
		}
		user = user[1:]
	}

	i := girc.New(girc.Config{
		Server:     server,
		ServerPass: b.GetString("Password"),
		Port:       port,
		Nick:       b.GetString("Nick"),
		User:       user,
		Name:       b.GetString("Nick"),
		SSL:        b.GetBool("UseTLS"),
		TLSConfig:  &tls.Config{InsecureSkipVerify: b.GetBool("SkipTLSVerify"), ServerName: server}, //nolint:gosec
		PingDelay:  time.Minute,
	})
	return i, nil
}

func (b *Birc) endNames(client *girc.Client, event girc.Event) {
	channel := event.Params[1]
	sort.Strings(b.names[channel])
	maxNamesPerPost := (300 / b.nicksPerRow()) * b.nicksPerRow()
	for len(b.names[channel]) > maxNamesPerPost {
		b.Remote <- config.Message{Username: b.Nick, Text: b.formatnicks(b.names[channel][0:maxNamesPerPost]),
			Channel: channel, Account: b.Account}
		b.names[channel] = b.names[channel][maxNamesPerPost:]
	}
	b.Remote <- config.Message{Username: b.Nick, Text: b.formatnicks(b.names[channel]),
		Channel: channel, Account: b.Account}
	b.names[channel] = nil
	b.i.Handlers.Clear(girc.RPL_NAMREPLY)
	b.i.Handlers.Clear(girc.RPL_ENDOFNAMES)
}

func (b *Birc) skipPrivMsg(event girc.Event) bool {
	// Our nick can be changed
	b.Nick = b.i.GetNick()

	// freenode doesn't send 001 as first reply
	if event.Command == "NOTICE" {
		return true
	}
	// don't forward queries to the bot
	if event.Params[0] == b.Nick {
		return true
	}
	// don't forward message from ourself
	if event.Source.Name == b.Nick {
		return true
	}
	return false
}

func (b *Birc) nicksPerRow() int {
	return 4
}

func (b *Birc) storeNames(client *girc.Client, event girc.Event) {
	channel := event.Params[2]
	b.names[channel] = append(
		b.names[channel],
		strings.Split(strings.TrimSpace(event.Last()), " ")...)
}

func (b *Birc) formatnicks(nicks []string) string {
	return strings.Join(nicks, ", ") + " currently on IRC"
}
<commit_msg>Be less lossy when throttling IRC messages (#1004)<commit_after>package birc

import (
	"crypto/tls"
	"fmt"
	"hash/crc32"
	"net"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/42wim/matterbridge/bridge"
	"github.com/42wim/matterbridge/bridge/config"
	"github.com/42wim/matterbridge/bridge/helper"
	"github.com/lrstanley/girc"

	// We need to import the 'data' package as an implicit dependency.
	// See: https://godoc.org/github.com/paulrosania/go-charset/charset
	_ "github.com/paulrosania/go-charset/data"
)

type Birc struct {
	i                                         *girc.Client
	Nick                                      string
	names                                     map[string][]string
	connected                                 chan error
	Local                                     chan config.Message // local queue for flood control
	FirstConnection, authDone                 bool
	MessageDelay, MessageQueue, MessageLength int

	*bridge.Config
}

func New(cfg *bridge.Config) bridge.Bridger {
	b := &Birc{}
	b.Config = cfg
	b.Nick = b.GetString("Nick")
	b.names = make(map[string][]string)
	b.connected = make(chan error)
	if b.GetInt("MessageDelay") == 0 {
		b.MessageDelay = 1300
	} else {
		b.MessageDelay = b.GetInt("MessageDelay")
	}
	if b.GetInt("MessageQueue") == 0 {
		b.MessageQueue = 30
	} else {
		b.MessageQueue = b.GetInt("MessageQueue")
	}
	if b.GetInt("MessageLength") == 0 {
		b.MessageLength = 400
	} else {
		b.MessageLength = b.GetInt("MessageLength")
	}
	b.FirstConnection = true
	return b
}

func (b *Birc) Command(msg *config.Message) string {
	if msg.Text == "!users" {
		b.i.Handlers.Add(girc.RPL_NAMREPLY, b.storeNames)
		b.i.Handlers.Add(girc.RPL_ENDOFNAMES, b.endNames)
		b.i.Cmd.SendRaw("NAMES " + msg.Channel) //nolint:errcheck
	}
	return ""
}

func (b *Birc) Connect() error {
	b.Local = make(chan config.Message, b.MessageQueue+10)
	b.Log.Infof("Connecting %s", b.GetString("Server"))

	i, err := b.getClient()
	if err != nil {
		return err
	}

	if b.GetBool("UseSASL") {
		i.Config.SASL = &girc.SASLPlain{
			User: b.GetString("NickServNick"),
			Pass: b.GetString("NickServPassword"),
		}
	}

	i.Handlers.Add(girc.RPL_WELCOME, b.handleNewConnection)
	i.Handlers.Add(girc.RPL_ENDOFMOTD, b.handleOtherAuth)
	i.Handlers.Add(girc.ERR_NOMOTD, b.handleOtherAuth)
	i.Handlers.Add(girc.ALL_EVENTS, b.handleOther)
	b.i = i

	go b.doConnect()

	err = <-b.connected
	if err != nil {
		return fmt.Errorf("connection failed %s", err)
	}
	b.Log.Info("Connection succeeded")
	b.FirstConnection = false
	if b.GetInt("DebugLevel") == 0 {
		i.Handlers.Clear(girc.ALL_EVENTS)
	}
	go b.doSend()
	return nil
}

func (b *Birc) Disconnect() error {
	b.i.Close()
	close(b.Local)
	return nil
}

func (b *Birc) JoinChannel(channel config.ChannelInfo) error {
	// need to check if we have nickserv auth done before joining channels
	for {
		if b.authDone {
			break
		}
		time.Sleep(time.Second)
	}
	if channel.Options.Key != "" {
		b.Log.Debugf("using key %s for channel %s", channel.Options.Key, channel.Name)
		b.i.Cmd.JoinKey(channel.Name, channel.Options.Key)
	} else {
		b.i.Cmd.Join(channel.Name)
	}
	return nil
}

func (b *Birc) Send(msg config.Message) (string, error) {
	// ignore delete messages
	if msg.Event == config.EventMsgDelete {
		return "", nil
	}

	b.Log.Debugf("=> Receiving %#v", msg)

	// we can be in between reconnects #385
	if !b.i.IsConnected() {
		b.Log.Error("Not connected to server, dropping message")
		return "", nil
	}

	// Execute a command
	if strings.HasPrefix(msg.Text, "!") {
		b.Command(&msg)
	}

	// convert to specified charset
	if err := b.handleCharset(&msg); err != nil {
		return "", err
	}

	// handle files, return if we're done here
	if ok := b.handleFiles(&msg); ok {
		return "", nil
	}

	var msgLines []string
	if b.GetBool("MessageSplit") {
		msgLines = helper.GetSubLines(msg.Text, b.MessageLength)
	} else {
		msgLines = helper.GetSubLines(msg.Text, 0)
	}
	for i := range msgLines {
		if len(b.Local) >= b.MessageQueue {
			b.Log.Debugf("flooding, dropping message (queue at %d)", len(b.Local))
			return "", nil
		}

		msg.Text = msgLines[i]
		b.Local <- msg
	}
	return "", nil
}

func (b *Birc) doConnect() {
	for {
		if err := b.i.Connect(); err != nil {
			b.Log.Errorf("disconnect: error: %s", err)
			if b.FirstConnection {
				b.connected <- err
				return
			}
		} else {
			b.Log.Info("disconnect: client requested quit")
		}
		b.Log.Info("reconnecting in 30 seconds...")
		time.Sleep(30 * time.Second)
		b.i.Handlers.Clear(girc.RPL_WELCOME)
		b.i.Handlers.Add(girc.RPL_WELCOME, func(client *girc.Client, event girc.Event) {
			b.Remote <- config.Message{Username: "system", Text: "rejoin", Channel: "", Account: b.Account, Event: config.EventRejoinChannels}
			// set our correct nick on reconnect if necessary
			b.Nick = event.Source.Name
		})
	}
}

func (b *Birc) doSend() {
	rate := time.Millisecond * time.Duration(b.MessageDelay)
	throttle := time.NewTicker(rate)
	for msg := range b.Local {
		<-throttle.C
		username := msg.Username
		if b.GetBool("Colornicks") {
			checksum := crc32.ChecksumIEEE([]byte(msg.Username))
			colorCode := checksum%14 + 2 // quick fix - prevent white or black color codes
			username = fmt.Sprintf("\x03%02d%s\x0F", colorCode, msg.Username)
		}
		if msg.Event == config.EventUserAction {
			b.i.Cmd.Action(msg.Channel, username+msg.Text)
		} else {
			b.Log.Debugf("Sending to channel %s", msg.Channel)
			b.i.Cmd.Message(msg.Channel, username+msg.Text)
		}
	}
}

// validateInput validates the server/port/nick configuration. Returns a *girc.Client if successful
func (b *Birc) getClient() (*girc.Client, error) {
	server, portstr, err := net.SplitHostPort(b.GetString("Server"))
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portstr)
	if err != nil {
		return nil, err
	}
	// fix strict user handling of girc
	user := b.GetString("Nick")
	for !girc.IsValidUser(user) {
		if len(user) == 1 || len(user) == 0 {
			user = "matterbridge"
			break
		}
		user = user[1:]
	}

	i := girc.New(girc.Config{
		Server:     server,
		ServerPass: b.GetString("Password"),
		Port:       port,
		Nick:       b.GetString("Nick"),
		User:       user,
		Name:       b.GetString("Nick"),
		SSL:        b.GetBool("UseTLS"),
		TLSConfig:  &tls.Config{InsecureSkipVerify: b.GetBool("SkipTLSVerify"), ServerName: server}, //nolint:gosec
		PingDelay:  time.Minute,
	})
	return i, nil
}

func (b *Birc) endNames(client *girc.Client, event girc.Event) {
	channel := event.Params[1]
	sort.Strings(b.names[channel])
	maxNamesPerPost := (300 / b.nicksPerRow()) * b.nicksPerRow()
	for len(b.names[channel]) > maxNamesPerPost {
		b.Remote <- config.Message{Username: b.Nick, Text: b.formatnicks(b.names[channel][0:maxNamesPerPost]),
			Channel: channel, Account: b.Account}
		b.names[channel] = b.names[channel][maxNamesPerPost:]
	}
	b.Remote <- config.Message{Username: b.Nick, Text: b.formatnicks(b.names[channel]),
		Channel: channel, Account: b.Account}
	b.names[channel] = nil
	b.i.Handlers.Clear(girc.RPL_NAMREPLY)
	b.i.Handlers.Clear(girc.RPL_ENDOFNAMES)
}

func (b *Birc) skipPrivMsg(event girc.Event) bool {
	// Our nick can be changed
	b.Nick = b.i.GetNick()

	// freenode doesn't send 001 as first reply
	if event.Command == "NOTICE" {
		return true
	}
	// don't forward queries to the bot
	if event.Params[0] == b.Nick {
		return true
	}
	// don't forward message from ourself
	if event.Source.Name == b.Nick {
		return true
	}
	return false
}

func (b *Birc) nicksPerRow() int {
	return 4
}

func (b *Birc) storeNames(client *girc.Client, event girc.Event) {
	channel := event.Params[2]
	b.names[channel] = append(
		b.names[channel],
		strings.Split(strings.TrimSpace(event.Last()), " ")...)
}

func (b *Birc) formatnicks(nicks []string) string {
	return strings.Join(nicks, ", ") + " currently on IRC"
}
<|endoftext|>