<commit_before>package trace

import (
	"context"
	"net"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/stretchr/testify/assert"
	"github.com/stripe/veneur/ssf"
)

const ε = .00002

func TestStartTrace(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	const expectedParent int64 = 0
	start := time.Now()
	trace := StartTrace(resource)
	end := time.Now()

	between := end.After(trace.Start) && trace.Start.After(start)

	assert.Equal(t, trace.TraceID, trace.SpanID)
	assert.Equal(t, trace.ParentID, expectedParent)
	assert.Equal(t, trace.Resource, resource)
	assert.True(t, between)
}

func TestRecord(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	const metricName = "veneur.trace.test"
	const serviceName = "veneur-test"
	Service = serviceName

	// arbitrary
	const BufferSize = 1087152

	traceAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	assert.NoError(t, err)
	serverConn, err := net.ListenUDP("udp", traceAddr)
	assert.NoError(t, err)

	err = serverConn.SetReadBuffer(BufferSize)
	assert.NoError(t, err)

	respChan := make(chan []byte)
	kill := make(chan struct{})

	go func() {
		buf := make([]byte, BufferSize)
		n, _, err := serverConn.ReadFrom(buf)
		assert.NoError(t, err)

		buf = buf[:n]
		respChan <- buf
	}()

	go func() {
		<-time.After(5 * time.Second)
		kill <- struct{}{}
	}()

	trace := StartTrace(resource)
	trace.Status = ssf.SSFSample_CRITICAL
	trace.error = true

	tags := map[string]string{
		"error.msg":   "an error occurred!",
		"error.type":  "type error interface",
		"error.stack": "insert\nlots\nof\nstuff",
		"resource":    resource,
		"name":        metricName,
	}

	trace.Record(metricName, tags)
	end := time.Now()

	select {
	case _ = <-kill:
		assert.Fail(t, "timed out waiting for socket read")
	case resp := <-respChan:
		// Because this is marshalled using protobuf,
		// we can't expect the representation to be immutable
		// and cannot test the marshalled payload directly
		sample := &ssf.SSFSpan{}
		err := proto.Unmarshal(resp, sample)

		assert.NoError(t, err)

		timestamp := time.Unix(sample.StartTimestamp/1e9, 0)

		assert.Equal(t, trace.Start.Unix(), timestamp.Unix())

		duration := sample.EndTimestamp - sample.StartTimestamp

		// We don't know the exact duration, but we can assert on the interval
		assert.True(t, duration > 0, "Expected positive trace duration")
		upperBound := end.Sub(trace.Start).Nanoseconds()
		assert.True(t, duration < upperBound, "Expected trace duration (%d) to be less than upper bound %d", duration, upperBound)

		for _, metric := range sample.Metrics {
			assert.InEpsilon(t, metric.SampleRate, 0.1, ε)
		}

		assertTagEquals(t, sample, "resource", resource)
		assertTagEquals(t, sample, "name", metricName)
		assert.Equal(t, true, sample.Error)
		assert.Equal(t, serviceName, sample.Service)
		assert.Equal(t, tags, sample.Tags)
	}

}

func TestAttach(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	ctx := context.Background()

	parent := ctx.Value(traceKey)
	assert.Nil(t, parent, "Expected not to find parent in context before attaching")

	trace := StartTrace(resource)
	ctx2 := trace.Attach(ctx)

	parent = ctx2.Value(traceKey).(*Trace)
	assert.NotNil(t, parent, "Expected not to find parent in context before attaching")
}

func TestSpanFromContext(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	trace := StartTrace(resource)

	ctx := trace.Attach(context.Background())
	child := SpanFromContext(ctx)
	// Test the *grandchild* so that we can ensure that
	// the parent ID is set independently of the trace ID
	ctx = child.Attach(context.Background())
	grandchild := SpanFromContext(ctx)

	assert.Equal(t, child.TraceID, trace.SpanID)
	assert.Equal(t, child.TraceID, trace.TraceID)
	assert.Equal(t, child.ParentID, trace.SpanID)
	assert.Equal(t, grandchild.ParentID, child.SpanID)
	assert.Equal(t, grandchild.TraceID, trace.SpanID)
}

func TestStartChildSpan(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	root := StartTrace(resource)
	child := StartChildSpan(root)
	grandchild := StartChildSpan(child)

	assert.Equal(t, resource, child.Resource)
	assert.Equal(t, resource, grandchild.Resource)

	assert.Equal(t, root.SpanID, root.TraceID)
	assert.Equal(t, root.SpanID, child.TraceID)
	assert.Equal(t, root.SpanID, grandchild.TraceID)

	assert.Equal(t, root.SpanID, child.ParentID)
	assert.Equal(t, child.SpanID, grandchild.ParentID)
}

// Test that a Trace is correctly able to generate
// its spanContext representation from the point of view
// of its children
func TestTraceContextAsParent(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	trace := StartTrace(resource)

	ctx := trace.contextAsParent()

	assert.Equal(t, trace.TraceID, ctx.TraceID())
	assert.Equal(t, trace.SpanID, ctx.ParentID())
	assert.Equal(t, trace.Resource, ctx.Resource())
}

func TestNameTag(t *testing.T) {
	const name = "my.name.tag"
	tracer := Tracer{}
	span := tracer.StartSpan("resource", NameTag(name)).(*Span)
	assert.Equal(t, 1, len(span.Tags))
	assert.Equal(t, name, span.Tags["name"])

}

type localError struct {
	message string
}

func (le localError) Error() string {
	return le.message
}

func TestError(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	const errorMessage = "some error happened"
	err := localError{errorMessage}

	root := StartTrace(resource)
	root.Error(err)

	assert.Equal(t, root.Status, ssf.SSFSample_CRITICAL)
	assert.Equal(t, len(root.Tags), 3)

	for k, v := range root.Tags {
		switch k {
		case errorMessageTag:
			assert.Equal(t, v, err.Error())
		case errorTypeTag:
			assert.Equal(t, v, "localError")
		case errorStackTag:
			assert.Equal(t, v, err.Error())
		}
	}

}

func TestStripPackageName(t *testing.T) {
	type testCase struct {
		Name     string
		fname    string
		expected string
	}

	cases := []testCase{
		{
			Name:     "Method",
			fname:    "github.com/stripe/veneur.(*Server).Flush",
			expected: "veneur.(*Server).Flush",
		},
		{
			Name:     "NestedPackageMethod",
			fname:    "github.com/stripe/veneur/trace.(*Tracer).StartSpan",
			expected: "trace.(*Tracer).StartSpan",
		},
		{
			// This shouldn't be valid, but we should at least ensure we don't
			// cause a runtime panic if it's passed
			Name:     "TrailingSlash",
			fname:    "github.com/",
			expected: "github.com/",
		},
	}

	for _, tc := range cases {
		t.Run(tc.Name, func(t *testing.T) {
			assert.Equal(t, stripPackageName(tc.fname), tc.expected)
		})
	}
}

func assertTagEquals(t *testing.T, sample *ssf.SSFSpan, name, value string) {
	assert.Equal(t, value, sample.Tags[name])
}
<commit_msg>Add test for creating root span from StartSpanFromContext<commit_after>package trace

import (
	"context"
	"net"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/stretchr/testify/assert"
	"github.com/stripe/veneur/ssf"
)

const ε = .00002

func TestStartTrace(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	const expectedParent int64 = 0
	start := time.Now()
	trace := StartTrace(resource)
	end := time.Now()

	between := end.After(trace.Start) && trace.Start.After(start)

	assert.Equal(t, trace.TraceID, trace.SpanID)
	assert.Equal(t, trace.ParentID, expectedParent)
	assert.Equal(t, trace.Resource, resource)
	assert.True(t, between)
}

func TestRecord(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	const metricName = "veneur.trace.test"
	const serviceName = "veneur-test"
	Service = serviceName

	// arbitrary
	const BufferSize = 1087152

	traceAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	assert.NoError(t, err)
	serverConn, err := net.ListenUDP("udp", traceAddr)
	assert.NoError(t, err)

	err = serverConn.SetReadBuffer(BufferSize)
	assert.NoError(t, err)

	respChan := make(chan []byte)
	kill := make(chan struct{})

	go func() {
		buf := make([]byte, BufferSize)
		n, _, err := serverConn.ReadFrom(buf)
		assert.NoError(t, err)

		buf = buf[:n]
		respChan <- buf
	}()

	go func() {
		<-time.After(5 * time.Second)
		kill <- struct{}{}
	}()

	trace := StartTrace(resource)
	trace.Status = ssf.SSFSample_CRITICAL
	trace.error = true

	tags := map[string]string{
		"error.msg":   "an error occurred!",
		"error.type":  "type error interface",
		"error.stack": "insert\nlots\nof\nstuff",
		"resource":    resource,
		"name":        metricName,
	}

	trace.Record(metricName, tags)
	end := time.Now()

	select {
	case _ = <-kill:
		assert.Fail(t, "timed out waiting for socket read")
	case resp := <-respChan:
		// Because this is marshalled using protobuf,
		// we can't expect the representation to be immutable
		// and cannot test the marshalled payload directly
		sample := &ssf.SSFSpan{}
		err := proto.Unmarshal(resp, sample)

		assert.NoError(t, err)

		timestamp := time.Unix(sample.StartTimestamp/1e9, 0)

		assert.Equal(t, trace.Start.Unix(), timestamp.Unix())

		duration := sample.EndTimestamp - sample.StartTimestamp

		// We don't know the exact duration, but we can assert on the interval
		assert.True(t, duration > 0, "Expected positive trace duration")
		upperBound := end.Sub(trace.Start).Nanoseconds()
		assert.True(t, duration < upperBound, "Expected trace duration (%d) to be less than upper bound %d", duration, upperBound)

		for _, metric := range sample.Metrics {
			assert.InEpsilon(t, metric.SampleRate, 0.1, ε)
		}

		assertTagEquals(t, sample, "resource", resource)
		assertTagEquals(t, sample, "name", metricName)
		assert.Equal(t, true, sample.Error)
		assert.Equal(t, serviceName, sample.Service)
		assert.Equal(t, tags, sample.Tags)
	}

}

func TestAttach(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	ctx := context.Background()

	parent := ctx.Value(traceKey)
	assert.Nil(t, parent, "Expected not to find parent in context before attaching")

	trace := StartTrace(resource)
	ctx2 := trace.Attach(ctx)

	parent = ctx2.Value(traceKey).(*Trace)
	assert.NotNil(t, parent, "Expected not to find parent in context before attaching")
}

func TestSpanFromContext(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	trace := StartTrace(resource)

	ctx := trace.Attach(context.Background())
	child := SpanFromContext(ctx)
	// Test the *grandchild* so that we can ensure that
	// the parent ID is set independently of the trace ID
	ctx = child.Attach(context.Background())
	grandchild := SpanFromContext(ctx)

	assert.Equal(t, child.TraceID, trace.SpanID)
	assert.Equal(t, child.TraceID, trace.TraceID)
	assert.Equal(t, child.ParentID, trace.SpanID)
	assert.Equal(t, grandchild.ParentID, child.SpanID)
	assert.Equal(t, grandchild.TraceID, trace.SpanID)
}

// StartSpanFromContext should create a brand-new root span
// if the context does not contain a span
func TestSpanFromContextNoParent(t *testing.T) {
	const resource = "example"
	ctx := context.Background()

	span, _ := StartSpanFromContext(ctx, resource)

	assert.Equal(t, span.TraceID, span.SpanID)
	assert.Equal(t, int64(0), span.ParentID)
}

func TestStartChildSpan(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	root := StartTrace(resource)
	child := StartChildSpan(root)
	grandchild := StartChildSpan(child)

	assert.Equal(t, resource, child.Resource)
	assert.Equal(t, resource, grandchild.Resource)

	assert.Equal(t, root.SpanID, root.TraceID)
	assert.Equal(t, root.SpanID, child.TraceID)
	assert.Equal(t, root.SpanID, grandchild.TraceID)

	assert.Equal(t, root.SpanID, child.ParentID)
	assert.Equal(t, child.SpanID, grandchild.ParentID)
}

// Test that a Trace is correctly able to generate
// its spanContext representation from the point of view
// of its children
func TestTraceContextAsParent(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	trace := StartTrace(resource)

	ctx := trace.contextAsParent()

	assert.Equal(t, trace.TraceID, ctx.TraceID())
	assert.Equal(t, trace.SpanID, ctx.ParentID())
	assert.Equal(t, trace.Resource, ctx.Resource())
}

func TestNameTag(t *testing.T) {
	const name = "my.name.tag"
	tracer := Tracer{}
	span := tracer.StartSpan("resource", NameTag(name)).(*Span)
	assert.Equal(t, 1, len(span.Tags))
	assert.Equal(t, name, span.Tags["name"])

}

type localError struct {
	message string
}

func (le localError) Error() string {
	return le.message
}

func TestError(t *testing.T) {
	const resource = "Robert'); DROP TABLE students;"
	const errorMessage = "some error happened"
	err := localError{errorMessage}

	root := StartTrace(resource)
	root.Error(err)

	assert.Equal(t, root.Status, ssf.SSFSample_CRITICAL)
	assert.Equal(t, len(root.Tags), 3)

	for k, v := range root.Tags {
		switch k {
		case errorMessageTag:
			assert.Equal(t, v, err.Error())
		case errorTypeTag:
			assert.Equal(t, v, "localError")
		case errorStackTag:
			assert.Equal(t, v, err.Error())
		}
	}

}

func TestStripPackageName(t *testing.T) {
	type testCase struct {
		Name     string
		fname    string
		expected string
	}

	cases := []testCase{
		{
			Name:     "Method",
			fname:    "github.com/stripe/veneur.(*Server).Flush",
			expected: "veneur.(*Server).Flush",
		},
		{
			Name:     "NestedPackageMethod",
			fname:    "github.com/stripe/veneur/trace.(*Tracer).StartSpan",
			expected: "trace.(*Tracer).StartSpan",
		},
		{
			// This shouldn't be valid, but we should at least ensure we don't
			// cause a runtime panic if it's passed
			Name:     "TrailingSlash",
			fname:    "github.com/",
			expected: "github.com/",
		},
	}

	for _, tc := range cases {
		t.Run(tc.Name, func(t *testing.T) {
			assert.Equal(t, stripPackageName(tc.fname), tc.expected)
		})
	}
}

func assertTagEquals(t *testing.T, sample *ssf.SSFSpan, name, value string) {
	assert.Equal(t, value, sample.Tags[name])
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/boltdb/bolt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

var (
	errNoMigrationsYet = errors.New("no migrations have been added yet")
	levelKey           = []byte("level")
)

// migrations is a list of all the migration functions.
// Their indexes are used to determine when they are to be applied.
//
// When adding a new migration, make sure to add it to this list.
var migrations = []func(db *bolt.DB) error{
	InitialMigration,
	MigrateOriginalColorPreferredColor,
	MigrateTournamentRunnerupStringPerson,
	MigrateMatchScoreOrderKillOrder,
	MigrateMatchCommitToRound,
}

var EmptyDBError = errors.New("drunkenfall: Empty DB.")

// Migrate is the main migration entrypoint
//
// When called, it will check the database to see what migrations have already
// been applied. If that is lower than the const TopVersion, all migrations up
// to that point will sequentially be applied.
func Migrate(db *bolt.DB) error {
	err := db.View(func(tx *bolt.Tx) error {
		if tx.Bucket(TournamentKey) == nil {
			return EmptyDBError
		}
		return nil
	})
	if err == EmptyDBError {
		log.Print("Empty db detected, nothing to migrate.")
		err := db.Update(func(tx *bolt.Tx) error {
			return setVersion(tx, len(migrations))
		})
		return err
	}
	version, err := getVersion(db)

	if err == errNoMigrationsYet {
		// No migrations have been done yet, so lets do the first one by adding
		// the migration bucket.
		version = 0
	} else if err != nil {
		// Something actually went wrong. Oh noes.
		return err
	}

	// If version is lower than the latest known version, it's time to apply the
	// migrations!
	if version < len(migrations) {
		if err := applyMigrations(db, version); err != nil {
			log.Print("Error: Migration application failed ;'(")
			return err
		}
	} else {
		log.Print("Database up to date.")
	}

	return nil
}

func applyMigrations(db *bolt.DB, version int) error {
	log.Printf(" --- Migrating %d -> %d:\n", version, len(migrations))
	if err := backup(db, version, "db-migration-backup/"); err != nil {
		return err
	}

	// Run the new migrations and the new migrations only
	new := migrations[version:]
	for x, migration := range new {
		log.Printf("Applying migration %d\n", x+len(new))
		if err := migration(db); err != nil {
			log.Print("Migration failure: ", err)
			return err
		}
	}

	log.Printf(" --- Migrations applied successfully. <3\n")
	return nil
}

func backup(db *bolt.DB, version int, path string) error {
	_ = os.Mkdir(path, 0755)

	fn := fmt.Sprintf(
		"%d_from-v%d-to-v%d.db",
		time.Now().UnixNano(),
		version,
		len(migrations),
	)
	dst := filepath.Join(path, fn)

	data, _ := ioutil.ReadFile(db.Path())
	if err := ioutil.WriteFile(dst, data, 0644); err != nil {
		return err
	}

	log.Printf("Backed up into %s\n", fn)
	return nil
}

func getVersion(db *bolt.DB) (int, error) {
	var version int
	err := db.View(func(tx *bolt.Tx) error {
		var err error
		b := tx.Bucket(MigrationKey)
		if b == nil {
			return errNoMigrationsYet
		}

		x := b.Get(levelKey)
		version, err = strconv.Atoi(string(x))
		return err
	})

	return version, err
}

func setVersion(tx *bolt.Tx, version int) error {
	b, err := tx.CreateBucketIfNotExists(MigrationKey)
	if err != nil {
		return err
	}

	out, _ := json.Marshal(version)
	err = b.Put(levelKey, out)
	return err
}

// InitialMigration adds the db_version bucket with the sole entry
func InitialMigration(db *bolt.DB) error {
	return db.Update(func(tx *bolt.Tx) error {
		return setVersion(tx, 1)
	})
}
<commit_msg>Fix gas security suggestion<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/boltdb/bolt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

var (
	errNoMigrationsYet = errors.New("no migrations have been added yet")
	levelKey           = []byte("level")
)

// migrations is a list of all the migration functions.
// Their indexes are used to determine when they are to be applied.
//
// When adding a new migration, make sure to add it to this list.
var migrations = []func(db *bolt.DB) error{
	InitialMigration,
	MigrateOriginalColorPreferredColor,
	MigrateTournamentRunnerupStringPerson,
	MigrateMatchScoreOrderKillOrder,
	MigrateMatchCommitToRound,
}

var EmptyDBError = errors.New("drunkenfall: Empty DB.")

// Migrate is the main migration entrypoint
//
// When called, it will check the database to see what migrations have already
// been applied. If that is lower than the const TopVersion, all migrations up
// to that point will sequentially be applied.
func Migrate(db *bolt.DB) error {
	err := db.View(func(tx *bolt.Tx) error {
		if tx.Bucket(TournamentKey) == nil {
			return EmptyDBError
		}
		return nil
	})
	if err == EmptyDBError {
		log.Print("Empty db detected, nothing to migrate.")
		err := db.Update(func(tx *bolt.Tx) error {
			return setVersion(tx, len(migrations))
		})
		return err
	}
	version, err := getVersion(db)

	if err == errNoMigrationsYet {
		// No migrations have been done yet, so lets do the first one by adding
		// the migration bucket.
		version = 0
	} else if err != nil {
		// Something actually went wrong. Oh noes.
		return err
	}

	// If version is lower than the latest known version, it's time to apply the
	// migrations!
	if version < len(migrations) {
		if err := applyMigrations(db, version); err != nil {
			log.Print("Error: Migration application failed ;'(")
			return err
		}
	} else {
		log.Print("Database up to date.")
	}

	return nil
}

func applyMigrations(db *bolt.DB, version int) error {
	log.Printf(" --- Migrating %d -> %d:\n", version, len(migrations))
	if err := backup(db, version, "db-migration-backup/"); err != nil {
		return err
	}

	// Run the new migrations and the new migrations only
	new := migrations[version:]
	for x, migration := range new {
		log.Printf("Applying migration %d\n", x+len(new))
		if err := migration(db); err != nil {
			log.Print("Migration failure: ", err)
			return err
		}
	}

	log.Printf(" --- Migrations applied successfully. <3\n")
	return nil
}

func backup(db *bolt.DB, version int, path string) error {
	_ = os.Mkdir(path, 0700)

	fn := fmt.Sprintf(
		"%d_from-v%d-to-v%d.db",
		time.Now().UnixNano(),
		version,
		len(migrations),
	)
	dst := filepath.Join(path, fn)

	data, _ := ioutil.ReadFile(db.Path())
	if err := ioutil.WriteFile(dst, data, 0644); err != nil {
		return err
	}

	log.Printf("Backed up into %s\n", fn)
	return nil
}

func getVersion(db *bolt.DB) (int, error) {
	var version int
	err := db.View(func(tx *bolt.Tx) error {
		var err error
		b := tx.Bucket(MigrationKey)
		if b == nil {
			return errNoMigrationsYet
		}

		x := b.Get(levelKey)
		version, err = strconv.Atoi(string(x))
		return err
	})

	return version, err
}

func setVersion(tx *bolt.Tx, version int) error {
	b, err := tx.CreateBucketIfNotExists(MigrationKey)
	if err != nil {
		return err
	}

	out, _ := json.Marshal(version)
	err = b.Put(levelKey, out)
	return err
}

// InitialMigration adds the db_version bucket with the sole entry
func InitialMigration(db *bolt.DB) error {
	return db.Update(func(tx *bolt.Tx) error {
		return setVersion(tx, 1)
	})
}
<|endoftext|>


<commit_before>package flannel

import (
	"context"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/rancher/k3s/pkg/agent/util"
	"github.com/rancher/k3s/pkg/daemons/config"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	v1 "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	cniConf = `{
  "name":"cbr0",
  "cniVersion":"0.3.1",
  "plugins":[
    {
      "type":"flannel",
      "delegate":{
        "hairpinMode":true,
        "forceAddress":true,
        "isDefaultGateway":true
      }
    },
    {
      "type":"portmap",
      "capabilities":{
        "portMappings":true
      }
    }
  ]
}
`

	flannelConf = `{
	"Network": "%CIDR%",
	"Backend": %backend%
}
`

	vxlanBackend = `{
	"Type": "vxlan"
}`

	ipsecBackend = `{
	"Type": "ipsec",
	"UDPEncap": true,
	"PSK": "%psk%"
}`

	wireguardBackend = `{
	"Type": "extension",
	"PreStartupCommand": "wg genkey | tee privatekey | wg pubkey",
	"PostStartupCommand": "export SUBNET_IP=$(echo $SUBNET | cut -d'/' -f 1); ip link del flannel.1 2>/dev/null; echo $PATH >&2; wg-add.sh flannel.1 && wg set flannel.1 listen-port 51820 private-key privatekey && ip addr add $SUBNET_IP/32 dev flannel.1 && ip link set flannel.1 up && ip route add $NETWORK dev flannel.1",
	"ShutdownCommand": "ip link del flannel.1",
	"SubnetAddCommand": "read PUBLICKEY; wg set flannel.1 peer $PUBLICKEY endpoint $PUBLIC_IP:51820 allowed-ips $SUBNET",
	"SubnetRemoveCommand": "read PUBLICKEY; wg set flannel.1 peer $PUBLICKEY remove"
}`
)

func Prepare(ctx context.Context, nodeConfig *config.Node) error {
	if err := createCNIConf(nodeConfig.AgentConfig.CNIConfDir); err != nil {
		return err
	}

	return createFlannelConf(nodeConfig)
}

func Run(ctx context.Context, nodeConfig *config.Node, nodes v1.NodeInterface) error {
	nodeName := nodeConfig.AgentConfig.NodeName

	for {
		node, err := nodes.Get(nodeName, metav1.GetOptions{})
		if err == nil && node.Spec.PodCIDR != "" {
			break
		}
		if err == nil {
			logrus.Infof("waiting for node %s CIDR not assigned yet", nodeName)
		} else {
			logrus.Infof("waiting for node %s: %v", nodeName, err)
		}
		time.Sleep(2 * time.Second)
	}

	go func() {
		err := flannel(ctx, nodeConfig.FlannelIface, nodeConfig.FlannelConf, nodeConfig.AgentConfig.KubeConfigKubelet)
		logrus.Fatalf("flannel exited: %v", err)
	}()

	return nil
}

func createCNIConf(dir string) error {
	if dir == "" {
		return nil
	}
	p := filepath.Join(dir, "10-flannel.conflist")
	return util.WriteFile(p, cniConf)
}

func createFlannelConf(nodeConfig *config.Node) error {
	if nodeConfig.FlannelConf == "" {
		return nil
	}
	if nodeConfig.FlannelConfOverride {
		logrus.Infof("Using custom flannel conf defined at %s", nodeConfig.FlannelConf)
		return nil
	}
	confJSON := strings.Replace(flannelConf, "%CIDR%", nodeConfig.AgentConfig.ClusterCIDR.String(), -1)

	var backendConf string

	switch nodeConfig.FlannelBackend {
	case config.FlannelBackendVXLAN:
		backendConf = vxlanBackend
	case config.FlannelBackendIPSEC:
		backendConf = strings.Replace(ipsecBackend, "%psk%", nodeConfig.AgentConfig.IPSECPSK, -1)
		if err := setupStrongSwan(nodeConfig); err != nil {
			return err
		}
	case config.FlannelBackendWireguard:
		backendConf = wireguardBackend
	default:
		return fmt.Errorf("Cannot configure unknown flannel backend '%s'", nodeConfig.FlannelBackend)
	}
	confJSON = strings.Replace(confJSON, "%backend%", backendConf, -1)

	return util.WriteFile(nodeConfig.FlannelConf, confJSON)
}

func setupStrongSwan(nodeConfig *config.Node) error {
	// if data dir env is not set point to root
	dataDir := os.Getenv("K3S_DATA_DIR")
	if dataDir == "" {
		dataDir = "/"
	}
	dataDir = path.Join(dataDir, "etc", "strongswan")

	info, err := os.Lstat(nodeConfig.AgentConfig.StrongSwanDir)
	// something exists but is not a symlink, return
	if err == nil && info.Mode()&os.ModeSymlink == 0 {
		return nil
	}
	if err == nil {
		target, err := os.Readlink(nodeConfig.AgentConfig.StrongSwanDir)
		// current link is the same, return
		if err == nil && target == dataDir {
			return nil
		}
	}

	// clean up strongswan old link
	os.Remove(nodeConfig.AgentConfig.StrongSwanDir)

	// make new strongswan link
	return os.Symlink(dataDir, nodeConfig.AgentConfig.StrongSwanDir)
}
<commit_msg>Set Wireguard keepalive to 25 seconds<commit_after>package flannel

import (
	"context"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/rancher/k3s/pkg/agent/util"
	"github.com/rancher/k3s/pkg/daemons/config"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	v1 "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	cniConf = `{
  "name":"cbr0",
  "cniVersion":"0.3.1",
  "plugins":[
    {
      "type":"flannel",
      "delegate":{
        "hairpinMode":true,
        "forceAddress":true,
        "isDefaultGateway":true
      }
    },
    {
      "type":"portmap",
      "capabilities":{
        "portMappings":true
      }
    }
  ]
}
`

	flannelConf = `{
	"Network": "%CIDR%",
	"Backend": %backend%
}
`

	vxlanBackend = `{
	"Type": "vxlan"
}`

	ipsecBackend = `{
	"Type": "ipsec",
	"UDPEncap": true,
	"PSK": "%psk%"
}`

	wireguardBackend = `{
	"Type": "extension",
	"PreStartupCommand": "wg genkey | tee privatekey | wg pubkey",
	"PostStartupCommand": "export SUBNET_IP=$(echo $SUBNET | cut -d'/' -f 1); ip link del flannel.1 2>/dev/null; echo $PATH >&2; wg-add.sh flannel.1 && wg set flannel.1 listen-port 51820 private-key privatekey persistent-keepalive 25 && ip addr add $SUBNET_IP/32 dev flannel.1 && ip link set flannel.1 up && ip route add $NETWORK dev flannel.1",
	"ShutdownCommand": "ip link del flannel.1",
	"SubnetAddCommand": "read PUBLICKEY; wg set flannel.1 peer $PUBLICKEY endpoint $PUBLIC_IP:51820 allowed-ips $SUBNET",
	"SubnetRemoveCommand": "read PUBLICKEY; wg set flannel.1 peer $PUBLICKEY remove"
}`
)

func Prepare(ctx context.Context, nodeConfig *config.Node) error {
	if err := createCNIConf(nodeConfig.AgentConfig.CNIConfDir); err != nil {
		return err
	}

	return createFlannelConf(nodeConfig)
}

func Run(ctx context.Context, nodeConfig *config.Node, nodes v1.NodeInterface) error {
	nodeName := nodeConfig.AgentConfig.NodeName

	for {
		node, err := nodes.Get(nodeName, metav1.GetOptions{})
		if err == nil && node.Spec.PodCIDR != "" {
			break
		}
		if err == nil {
			logrus.Infof("waiting for node %s CIDR not assigned yet", nodeName)
		} else {
			logrus.Infof("waiting for node %s: %v", nodeName, err)
		}
		time.Sleep(2 * time.Second)
	}

	go func() {
		err := flannel(ctx, nodeConfig.FlannelIface, nodeConfig.FlannelConf, nodeConfig.AgentConfig.KubeConfigKubelet)
		logrus.Fatalf("flannel exited: %v", err)
	}()

	return nil
}

func createCNIConf(dir string) error {
	if dir == "" {
		return nil
	}
	p := filepath.Join(dir, "10-flannel.conflist")
	return util.WriteFile(p, cniConf)
}

func createFlannelConf(nodeConfig *config.Node) error {
	if nodeConfig.FlannelConf == "" {
		return nil
	}
	if nodeConfig.FlannelConfOverride {
		logrus.Infof("Using custom flannel conf defined at %s", nodeConfig.FlannelConf)
		return nil
	}
	confJSON := strings.Replace(flannelConf, "%CIDR%", nodeConfig.AgentConfig.ClusterCIDR.String(), -1)

	var backendConf string

	switch nodeConfig.FlannelBackend {
	case config.FlannelBackendVXLAN:
		backendConf = vxlanBackend
	case config.FlannelBackendIPSEC:
		backendConf = strings.Replace(ipsecBackend, "%psk%", nodeConfig.AgentConfig.IPSECPSK, -1)
		if err := setupStrongSwan(nodeConfig); err != nil {
			return err
		}
	case config.FlannelBackendWireguard:
		backendConf = wireguardBackend
	default:
		return fmt.Errorf("Cannot configure unknown flannel backend '%s'", nodeConfig.FlannelBackend)
	}
	confJSON = strings.Replace(confJSON, "%backend%", backendConf, -1)

	return util.WriteFile(nodeConfig.FlannelConf, confJSON)
}

func setupStrongSwan(nodeConfig *config.Node) error {
	// if data dir env is not set point to root
	dataDir := os.Getenv("K3S_DATA_DIR")
	if dataDir == "" {
		dataDir = "/"
	}
	dataDir = path.Join(dataDir, "etc", "strongswan")

	info, err := os.Lstat(nodeConfig.AgentConfig.StrongSwanDir)
	// something exists but is not a symlink, return
	if err == nil && info.Mode()&os.ModeSymlink == 0 {
		return nil
	}
	if err == nil {
		target, err := os.Readlink(nodeConfig.AgentConfig.StrongSwanDir)
		// current link is the same, return
		if err == nil && target == dataDir {
			return nil
		}
	}

	// clean up strongswan old link
	os.Remove(nodeConfig.AgentConfig.StrongSwanDir)

	// make new strongswan link
	return os.Symlink(dataDir, nodeConfig.AgentConfig.StrongSwanDir)
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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 api_test

import (
	"io/ioutil"
	"testing"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/testapi"
	_ "github.com/GoogleCloudPlatform/kubernetes/pkg/api/v1beta3"
)

func BenchmarkPodConversion(b *testing.B) {
	data, err := ioutil.ReadFile("pod_example.json")
	if err != nil {
		b.Fatalf("Unexpected error while reading file: %v", err)
	}
	var pod api.Pod
	if err := api.Scheme.DecodeInto(data, &pod); err != nil {
		b.Fatalf("Unexpected error decoding pod: %v", err)
	}

	scheme := api.Scheme.Raw()
	var result *api.Pod
	for i := 0; i < b.N; i++ {
		versionedObj, err := scheme.ConvertToVersion(&pod, testapi.Version())
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		obj, err := scheme.ConvertToVersion(versionedObj, scheme.InternalVersion)
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		result = obj.(*api.Pod)
	}
	if !api.Semantic.DeepDerivative(pod, *result) {
		b.Fatalf("Incorrect conversion: expected %v, got %v", pod, *result)
	}
}

func BenchmarkNodeConversion(b *testing.B) {
	data, err := ioutil.ReadFile("node_example.json")
	if err != nil {
		b.Fatalf("Unexpected error while reading file: %v", err)
	}
	var node api.Node
	if err := api.Scheme.DecodeInto(data, &node); err != nil {
		b.Fatalf("Unexpected error decoding node: %v", err)
	}

	scheme := api.Scheme.Raw()
	var result *api.Node
	for i := 0; i < b.N; i++ {
		versionedObj, err := scheme.ConvertToVersion(&node, testapi.Version())
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		obj, err := scheme.ConvertToVersion(versionedObj, scheme.InternalVersion)
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		result = obj.(*api.Node)
	}
	if !api.Semantic.DeepDerivative(node, *result) {
		b.Fatalf("Incorrect conversion: expected %v, got %v", node, *result)
	}
}

func BenchmarkReplicationControllerConversion(b *testing.B) {
	data, err := ioutil.ReadFile("replication_controller_example.json")
	if err != nil {
		b.Fatalf("Unexpected error while reading file: %v", err)
	}
	var replicationController api.ReplicationController
	if err := api.Scheme.DecodeInto(data, &replicationController); err != nil {
		b.Fatalf("Unexpected error decoding node: %v", err)
	}

	scheme := api.Scheme.Raw()
	var result *api.ReplicationController
	for i := 0; i < b.N; i++ {
		versionedObj, err := scheme.ConvertToVersion(&replicationController, testapi.Version())
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		obj, err := scheme.ConvertToVersion(versionedObj, scheme.InternalVersion)
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		result = obj.(*api.ReplicationController)
	}
	if !api.Semantic.DeepDerivative(replicationController, *result) {
		b.Fatalf("Incorrect conversion: expected %v, got %v", replicationController, *result)
	}
}
<commit_msg>remove import of v1beta3 package in api/conversion_test.go<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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 api_test

import (
	"io/ioutil"
	"testing"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/testapi"
)

func BenchmarkPodConversion(b *testing.B) {
	data, err := ioutil.ReadFile("pod_example.json")
	if err != nil {
		b.Fatalf("Unexpected error while reading file: %v", err)
	}
	var pod api.Pod
	if err := api.Scheme.DecodeInto(data, &pod); err != nil {
		b.Fatalf("Unexpected error decoding pod: %v", err)
	}

	scheme := api.Scheme.Raw()
	var result *api.Pod
	for i := 0; i < b.N; i++ {
		versionedObj, err := scheme.ConvertToVersion(&pod, testapi.Version())
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		obj, err := scheme.ConvertToVersion(versionedObj, scheme.InternalVersion)
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		result = obj.(*api.Pod)
	}
	if !api.Semantic.DeepDerivative(pod, *result) {
		b.Fatalf("Incorrect conversion: expected %v, got %v", pod, *result)
	}
}

func BenchmarkNodeConversion(b *testing.B) {
	data, err := ioutil.ReadFile("node_example.json")
	if err != nil {
		b.Fatalf("Unexpected error while reading file: %v", err)
	}
	var node api.Node
	if err := api.Scheme.DecodeInto(data, &node); err != nil {
		b.Fatalf("Unexpected error decoding node: %v", err)
	}

	scheme := api.Scheme.Raw()
	var result *api.Node
	for i := 0; i < b.N; i++ {
		versionedObj, err := scheme.ConvertToVersion(&node, testapi.Version())
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		obj, err := scheme.ConvertToVersion(versionedObj, scheme.InternalVersion)
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		result = obj.(*api.Node)
	}
	if !api.Semantic.DeepDerivative(node, *result) {
		b.Fatalf("Incorrect conversion: expected %v, got %v", node, *result)
	}
}

func BenchmarkReplicationControllerConversion(b *testing.B) {
	data, err := ioutil.ReadFile("replication_controller_example.json")
	if err != nil {
		b.Fatalf("Unexpected error while reading file: %v", err)
	}
	var replicationController api.ReplicationController
	if err := api.Scheme.DecodeInto(data, &replicationController); err != nil {
		b.Fatalf("Unexpected error decoding node: %v", err)
	}

	scheme := api.Scheme.Raw()
	var result *api.ReplicationController
	for i := 0; i < b.N; i++ {
		versionedObj, err := scheme.ConvertToVersion(&replicationController, testapi.Version())
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		obj, err := scheme.ConvertToVersion(versionedObj, scheme.InternalVersion)
		if err != nil {
			b.Fatalf("Conversion error: %v", err)
		}
		result = obj.(*api.ReplicationController)
	}
	if !api.Semantic.DeepDerivative(replicationController, *result) {
		b.Fatalf("Incorrect conversion: expected %v, got %v", replicationController, *result)
	}
}
<|endoftext|>


<commit_before>package wormhole

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"

	"github.com/golang/glog"
	"github.com/koding/tunnel"
	"github.com/sapcc/kubernikus/pkg/cmd"
	"github.com/sapcc/kubernikus/pkg/log"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/clientcmd/api"
)

func NewClientCommand() *cobra.Command {
	o := NewClientOptions()

	c := &cobra.Command{
		Use:   "client",
		Short: "Creates a Wormhole Client",
		Run: func(c *cobra.Command, args []string) {
			cmd.CheckError(o.Validate(c, args))
			cmd.CheckError(o.Complete(args))
			cmd.CheckError(o.Run(c))
		},
	}

	o.BindFlags(c.Flags())

	return c
}

type ClientOptions struct {
	KubeConfig string
	Server     string
	Context    string
}

func NewClientOptions() *ClientOptions {
	return &ClientOptions{}
}

func (o *ClientOptions) BindFlags(flags *pflag.FlagSet) {
	flags.StringVar(&o.KubeConfig, "kubeconfig", o.KubeConfig, "Path to the kubeconfig file to use to talk to the Tunnel Server. If unset, try the environment variable KUBECONFIG, as well as in-cluster configuration")
	flags.StringVar(&o.Server, "server", o.Server, "Tunnel Server endpoint (host:port)")
	flags.StringVar(&o.Context, "context", o.Context, "Kubeconfig context to use. (default: current-context)")
}

func (o *ClientOptions) Validate(c *cobra.Command, args []string) error {
	return nil
}

func (o *ClientOptions) Complete(args []string) error {
	return nil
}

func (o *ClientOptions) Run(c *cobra.Command) error {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, syscall.SIGTERM) // Push signals into channel
	wg := &sync.WaitGroup{}

	config, err := clientcmd.LoadFromFile(o.KubeConfig)
	if err != nil {
		return fmt.Errorf("Failed to load kubeconfig file %#v: %s", o.KubeConfig, err)
	}
	err = api.FlattenConfig(config)
	if err != nil {
		return err
	}
	contextName := config.CurrentContext
	if contextName == "" {
		contextName = o.Context
	}
	if contextName == "" {
		return fmt.Errorf("No context given")
	}

	context, found := config.Contexts[contextName]
	if !found {
		return fmt.Errorf("Context %s not found", contextName)
	}

	cluster, found := config.Clusters[context.Cluster]
	if !found {
		return fmt.Errorf("Cluster not found %s", context.Cluster)
	}

	authInfo, found := config.AuthInfos[context.AuthInfo]
	if !found {
		return fmt.Errorf("No auth info found for context %s", context.AuthInfo)
	}
	cert := authInfo.ClientCertificateData
	key := authInfo.ClientKeyData

	ca := cluster.CertificateAuthorityData

	var rootCAs *x509.CertPool
	if ca != nil {
		rootCAs = x509.NewCertPool()
		if !rootCAs.AppendCertsFromPEM(ca) {
			return fmt.Errorf("Failed to load any certs from %s", ca)
		}
	}
	certificate, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return fmt.Errorf("Failed to load certificate/key: %s", err)
	}

	x509cert, err := x509.ParseCertificate(certificate.Certificate[0])
	if err != nil {
		return fmt.Errorf("Failed to extract common name from client cert: %s", err)
	}

	serverAddr := o.Server
	if serverAddr == "" {
		url, err := url.Parse(cluster.Server)
		if err != nil {
			return err
		}
		c := strings.Split(url.Hostname(), ".")
		//Add "-t" to first component of hostname
		c[0] = fmt.Sprintf("%s-wormhole", c[0])
		serverAddr = fmt.Sprintf("%s:%s", strings.Join(c, "."), url.Port())
	}

	cfg := &tunnel.ClientConfig{
		Log:        new(log.KodingToGlogAdapter),
		Debug:      true,
		Identifier: x509cert.Subject.CommonName,
		ServerAddr: serverAddr,
		Dial: func(network, address string) (net.Conn, error) {
			conn, err := tls.Dial(network, address, &tls.Config{
				RootCAs:      rootCAs,
				Certificates: []tls.Certificate{certificate},
			})
			if err != nil {
				glog.Warningf("Failed to connect to %s: %s", address, err)
			}
			return conn, err
		},
	}

	client, err := tunnel.NewClient(cfg)
	if err != nil {
		return err
	}
	glog.Infof("Connecting to %s with id %s", cfg.ServerAddr, cfg.Identifier)

	wg.Add(1)
	go func() {
		defer wg.Done()
		client.Start()
	}()

	<-sigs
	glog.Info("Shutting down...")
	client.Close()
	wg.Wait()

	return nil
}
<commit_msg>url port is empty. hard-coded to ssl for now<commit_after>package wormhole

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"

	"github.com/golang/glog"
	"github.com/koding/tunnel"
	"github.com/sapcc/kubernikus/pkg/cmd"
	"github.com/sapcc/kubernikus/pkg/log"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/clientcmd/api"
)

func NewClientCommand() *cobra.Command {
	o := NewClientOptions()

	c := &cobra.Command{
		Use:   "client",
		Short: "Creates a Wormhole Client",
		Run: func(c *cobra.Command, args []string) {
			cmd.CheckError(o.Validate(c, args))
			cmd.CheckError(o.Complete(args))
			cmd.CheckError(o.Run(c))
		},
	}

	o.BindFlags(c.Flags())

	return c
}

type ClientOptions struct {
	KubeConfig string
	Server     string
	Context    string
}

func NewClientOptions() *ClientOptions {
	return &ClientOptions{}
}

func (o *ClientOptions) BindFlags(flags *pflag.FlagSet) {
	flags.StringVar(&o.KubeConfig, "kubeconfig", o.KubeConfig, "Path to the kubeconfig file to use to talk to the Tunnel Server. If unset, try the environment variable KUBECONFIG, as well as in-cluster configuration")
	flags.StringVar(&o.Server, "server", o.Server, "Tunnel Server endpoint (host:port)")
	flags.StringVar(&o.Context, "context", o.Context, "Kubeconfig context to use. (default: current-context)")
}

func (o *ClientOptions) Validate(c *cobra.Command, args []string) error {
	return nil
}

func (o *ClientOptions) Complete(args []string) error {
	return nil
}

func (o *ClientOptions) Run(c *cobra.Command) error {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, syscall.SIGTERM) // Push signals into channel
	wg := &sync.WaitGroup{}

	config, err := clientcmd.LoadFromFile(o.KubeConfig)
	if err != nil {
		return fmt.Errorf("Failed to load kubeconfig file %#v: %s", o.KubeConfig, err)
	}
	err = api.FlattenConfig(config)
	if err != nil {
		return err
	}
	contextName := config.CurrentContext
	if contextName == "" {
		contextName = o.Context
	}
	if contextName == "" {
		return fmt.Errorf("No context given")
	}

	context, found := config.Contexts[contextName]
	if !found {
		return fmt.Errorf("Context %s not found", contextName)
	}

	cluster, found := config.Clusters[context.Cluster]
	if !found {
		return fmt.Errorf("Cluster not found %s", context.Cluster)
	}

	authInfo, found := config.AuthInfos[context.AuthInfo]
	if !found {
		return fmt.Errorf("No auth info found for context %s", context.AuthInfo)
	}
	cert := authInfo.ClientCertificateData
	key := authInfo.ClientKeyData

	ca := cluster.CertificateAuthorityData

	var rootCAs *x509.CertPool
	if ca != nil {
		rootCAs = x509.NewCertPool()
		if !rootCAs.AppendCertsFromPEM(ca) {
			return fmt.Errorf("Failed to load any certs from %s", ca)
		}
	}
	certificate, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return fmt.Errorf("Failed to load certificate/key: %s", err)
	}

	x509cert, err := x509.ParseCertificate(certificate.Certificate[0])
	if err != nil {
		return fmt.Errorf("Failed to extract common name from client cert: %s", err)
	}

	serverAddr := o.Server
	if serverAddr == "" {
		url, err := url.Parse(cluster.Server)
		if err != nil {
			return err
		}
		c := strings.Split(url.Hostname(), ".")
		//Add "-t" to first component of hostname
		c[0] = fmt.Sprintf("%s-wormhole", c[0])
		serverAddr = fmt.Sprintf("%s:%s", strings.Join(c, "."), "443")
	}

	cfg := &tunnel.ClientConfig{
		Log:        new(log.KodingToGlogAdapter),
		Debug:      true,
		Identifier: x509cert.Subject.CommonName,
		ServerAddr: serverAddr,
		Dial: func(network, address string) (net.Conn, error) {
			conn, err := tls.Dial(network, address, &tls.Config{
				RootCAs:      rootCAs,
				Certificates: []tls.Certificate{certificate},
			})
			if err != nil {
				glog.Warningf("Failed to connect to %s: %s", address, err)
			}
			return conn, err
		},
	}

	client, err := tunnel.NewClient(cfg)
	if err != nil {
		return err
	}
	glog.Infof("Connecting to %s with id %s", cfg.ServerAddr, cfg.Identifier)

	wg.Add(1)
	go func() {
		defer wg.Done()
		client.Start()
	}()

	<-sigs
	glog.Info("Shutting down...")
	client.Close()
	wg.Wait()

	return nil
}
<|endoftext|>


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

package encrypt

import (
	"crypto/md5"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"

	"golang.org/x/crypto/argon2"
)

const (
	// sseGenericHeader is the AWS SSE header used for SSE-S3 and SSE-KMS.
	sseGenericHeader = "X-Amz-Server-Side-Encryption"

	// sseKmsKeyID is the AWS SSE-KMS key id.
	sseKmsKeyID = sseGenericHeader + "-Aws-Kms-Key-Id"
	// sseEncryptionContext is the AWS SSE-KMS Encryption Context data.
	sseEncryptionContext = sseGenericHeader + "-Encryption-Context"

	// sseCustomerAlgorithm is the AWS SSE-C algorithm HTTP header key.
	sseCustomerAlgorithm = sseGenericHeader + "-Customer-Algorithm"
	// sseCustomerKey is the AWS SSE-C encryption key HTTP header key.
	sseCustomerKey = sseGenericHeader + "-Customer-Key"
	// sseCustomerKeyMD5 is the AWS SSE-C encryption key MD5 HTTP header key.
	sseCustomerKeyMD5 = sseGenericHeader + "-Customer-Key-MD5"

	// sseCopyCustomerAlgorithm is the AWS SSE-C algorithm HTTP header key for CopyObject API.
	sseCopyCustomerAlgorithm = "X-Amz-Copy-Source-Server-Side-Encryption-Customer-Algorithm"
	// sseCopyCustomerKey is the AWS SSE-C encryption key HTTP header key for CopyObject API.
	sseCopyCustomerKey = "X-Amz-Copy-Source-Server-Side-Encryption-Customer-Key"
	// sseCopyCustomerKeyMD5 is the AWS SSE-C encryption key MD5 HTTP header key for CopyObject API.
	sseCopyCustomerKeyMD5 = "X-Amz-Copy-Source-Server-Side-Encryption-Customer-Key-MD5"
)

// PBKDF creates a SSE-C key from the provided password and salt.
// PBKDF is a password-based key derivation function
// which can be used to derive a high-entropy cryptographic
// key from a low-entropy password and a salt.
type PBKDF func(password, salt []byte) ServerSide

// DefaultPBKDF is the default PBKDF. It uses Argon2id with the
// recommended parameters from the RFC draft (1 pass, 64 MB memory, 4 threads).
var DefaultPBKDF PBKDF = func(password, salt []byte) ServerSide {
	sse := ssec{}
	copy(sse[:], argon2.IDKey(password, salt, 1, 64*1024, 4, 32))
	return sse
}

// Type is the server-side-encryption method. It represents one of
// the following encryption methods:
//  - SSE-C: server-side-encryption with customer provided keys
//  - KMS:   server-side-encryption with managed keys
//  - S3:    server-side-encryption using S3 storage encryption
type Type string

const (
	// SSEC represents server-side-encryption with customer provided keys
	SSEC Type = "SSE-C"
	// KMS represents server-side-encryption with managed keys
	KMS Type = "KMS"
	// S3 represents server-side-encryption using S3 storage encryption
	S3 Type = "S3"
)

// ServerSide is a form of S3 server-side-encryption.
type ServerSide interface {
	// Type returns the server-side-encryption method.
	Type() Type

	// Marshal adds encryption headers to the provided HTTP headers.
	// It marks an HTTP request as server-side-encryption request
	// and inserts the required data into the headers.
	Marshal(h http.Header)
}

// NewSSE returns a server-side-encryption using S3 storage encryption.
// Using SSE-S3 the server will encrypt the object with server-managed keys.
func NewSSE() ServerSide { return s3{} }

// NewSSEKMS returns a new server-side-encryption using SSE-KMS and the provided Key Id and context.
func NewSSEKMS(keyID string, context interface{}) (ServerSide, error) {
	if context == nil {
		return kms{key: keyID, hasContext: false}, nil
	}
	serializedContext, err := json.Marshal(context)
	if err != nil {
		return nil, err
	}
	return kms{key: keyID, context: serializedContext, hasContext: true}, nil
}

// NewSSEC returns a new server-side-encryption using SSE-C and the provided key.
// The key must be 32 bytes long.
func NewSSEC(key []byte) (ServerSide, error) {
	if len(key) != 32 {
		return nil, errors.New("encrypt: SSE-C key must be 256 bit long")
	}
	sse := ssec{}
	copy(sse[:], key)
	return sse, nil
}

// SSE transforms a SSE-C copy encryption into a SSE-C encryption.
// It is the inverse of SSECopy(...).
//
// If the provided sse is no SSE-C copy encryption SSE returns
// sse unmodified.
func SSE(sse ServerSide) ServerSide {
	if sse == nil || sse.Type() != SSEC {
		return sse
	}
	if sse, ok := sse.(ssecCopy); ok {
		return ssec(sse)
	}
	return sse
}

// SSECopy transforms a SSE-C encryption into a SSE-C copy
// encryption. This is required for SSE-C key rotation or a SSE-C
// copy where the source and the destination should be encrypted.
//
// If the provided sse is no SSE-C encryption SSECopy returns
// sse unmodified.
func SSECopy(sse ServerSide) ServerSide {
	if sse == nil || sse.Type() != SSEC {
		return sse
	}
	if sse, ok := sse.(ssec); ok {
		return ssecCopy(sse)
	}
	return sse
}

type ssec [32]byte

func (s ssec) Type() Type { return SSEC }

func (s ssec) Marshal(h http.Header) {
	keyMD5 := md5.Sum(s[:])
	h.Set(sseCustomerAlgorithm, "AES256")
	h.Set(sseCustomerKey, base64.StdEncoding.EncodeToString(s[:]))
	h.Set(sseCustomerKeyMD5, base64.StdEncoding.EncodeToString(keyMD5[:]))
}

type ssecCopy [32]byte

func (s ssecCopy) Type() Type { return SSEC }

func (s ssecCopy) Marshal(h http.Header) {
	keyMD5 := md5.Sum(s[:])
	h.Set(sseCopyCustomerAlgorithm, "AES256")
	h.Set(sseCopyCustomerKey, base64.StdEncoding.EncodeToString(s[:]))
	h.Set(sseCopyCustomerKeyMD5, base64.StdEncoding.EncodeToString(keyMD5[:]))
}

type s3 struct{}

func (s s3) Type() Type { return S3 }

func (s s3) Marshal(h http.Header) { h.Set(sseGenericHeader, "AES256") }

type kms struct {
	key        string
	context    []byte
	hasContext bool
}

func (s kms) Type() Type { return KMS }

func (s kms) Marshal(h http.Header) {
	h.Set(sseGenericHeader, "aws:kms")
	h.Set(sseKmsKeyID, s.key)
	if s.hasContext {
		h.Set(sseEncryptionContext, base64.StdEncoding.EncodeToString(s.context))
	}
}
<commit_msg>Remove empty X-Amz-Server-Side-Encryption-Aws-Kms-Key-Id request header (#1216)<commit_after>/*
 * MinIO Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2018 MinIO, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package encrypt

import (
	"crypto/md5"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"

	"golang.org/x/crypto/argon2"
)

const (
	// sseGenericHeader is the AWS SSE header used for SSE-S3 and SSE-KMS.
	sseGenericHeader = "X-Amz-Server-Side-Encryption"

	// sseKmsKeyID is the AWS SSE-KMS key id.
	sseKmsKeyID = sseGenericHeader + "-Aws-Kms-Key-Id"
	// sseEncryptionContext is the AWS SSE-KMS Encryption Context data.
	sseEncryptionContext = sseGenericHeader + "-Encryption-Context"

	// sseCustomerAlgorithm is the AWS SSE-C algorithm HTTP header key.
	sseCustomerAlgorithm = sseGenericHeader + "-Customer-Algorithm"
	// sseCustomerKey is the AWS SSE-C encryption key HTTP header key.
	sseCustomerKey = sseGenericHeader + "-Customer-Key"
	// sseCustomerKeyMD5 is the AWS SSE-C encryption key MD5 HTTP header key.
	sseCustomerKeyMD5 = sseGenericHeader + "-Customer-Key-MD5"

	// sseCopyCustomerAlgorithm is the AWS SSE-C algorithm HTTP header key for CopyObject API.
	sseCopyCustomerAlgorithm = "X-Amz-Copy-Source-Server-Side-Encryption-Customer-Algorithm"
	// sseCopyCustomerKey is the AWS SSE-C encryption key HTTP header key for CopyObject API.
	sseCopyCustomerKey = "X-Amz-Copy-Source-Server-Side-Encryption-Customer-Key"
	// sseCopyCustomerKeyMD5 is the AWS SSE-C encryption key MD5 HTTP header key for CopyObject API.
	sseCopyCustomerKeyMD5 = "X-Amz-Copy-Source-Server-Side-Encryption-Customer-Key-MD5"
)

// PBKDF creates a SSE-C key from the provided password and salt.
// PBKDF is a password-based key derivation function
// which can be used to derive a high-entropy cryptographic
// key from a low-entropy password and a salt.
type PBKDF func(password, salt []byte) ServerSide

// DefaultPBKDF is the default PBKDF. It uses Argon2id with the
// recommended parameters from the RFC draft (1 pass, 64 MB memory, 4 threads).
var DefaultPBKDF PBKDF = func(password, salt []byte) ServerSide {
	sse := ssec{}
	copy(sse[:], argon2.IDKey(password, salt, 1, 64*1024, 4, 32))
	return sse
}

// Type is the server-side-encryption method. It represents one of
// the following encryption methods:
//  - SSE-C: server-side-encryption with customer provided keys
//  - KMS:   server-side-encryption with managed keys
//  - S3:    server-side-encryption using S3 storage encryption
type Type string

const (
	// SSEC represents server-side-encryption with customer provided keys
	SSEC Type = "SSE-C"
	// KMS represents server-side-encryption with managed keys
	KMS Type = "KMS"
	// S3 represents server-side-encryption using S3 storage encryption
	S3 Type = "S3"
)

// ServerSide is a form of S3 server-side-encryption.
type ServerSide interface {
	// Type returns the server-side-encryption method.
	Type() Type

	// Marshal adds encryption headers to the provided HTTP headers.
	// It marks an HTTP request as server-side-encryption request
	// and inserts the required data into the headers.
	Marshal(h http.Header)
}

// NewSSE returns a server-side-encryption using S3 storage encryption.
// Using SSE-S3 the server will encrypt the object with server-managed keys.
func NewSSE() ServerSide { return s3{} }

// NewSSEKMS returns a new server-side-encryption using SSE-KMS and the provided Key Id and context.
func NewSSEKMS(keyID string, context interface{}) (ServerSide, error) {
	if context == nil {
		return kms{key: keyID, hasContext: false}, nil
	}
	serializedContext, err := json.Marshal(context)
	if err != nil {
		return nil, err
	}
	return kms{key: keyID, context: serializedContext, hasContext: true}, nil
}

// NewSSEC returns a new server-side-encryption using SSE-C and the provided key.
// The key must be 32 bytes long.
func NewSSEC(key []byte) (ServerSide, error) {
	if len(key) != 32 {
		return nil, errors.New("encrypt: SSE-C key must be 256 bit long")
	}
	sse := ssec{}
	copy(sse[:], key)
	return sse, nil
}

// SSE transforms a SSE-C copy encryption into a SSE-C encryption.
// It is the inverse of SSECopy(...).
//
// If the provided sse is no SSE-C copy encryption SSE returns
// sse unmodified.
func SSE(sse ServerSide) ServerSide {
	if sse == nil || sse.Type() != SSEC {
		return sse
	}
	if sse, ok := sse.(ssecCopy); ok {
		return ssec(sse)
	}
	return sse
}

// SSECopy transforms a SSE-C encryption into a SSE-C copy
// encryption. This is required for SSE-C key rotation or a SSE-C
// copy where the source and the destination should be encrypted.
//
// If the provided sse is no SSE-C encryption SSECopy returns
// sse unmodified.
func SSECopy(sse ServerSide) ServerSide {
	if sse == nil || sse.Type() != SSEC {
		return sse
	}
	if sse, ok := sse.(ssec); ok {
		return ssecCopy(sse)
	}
	return sse
}

type ssec [32]byte

func (s ssec) Type() Type { return SSEC }

func (s ssec) Marshal(h http.Header) {
	keyMD5 := md5.Sum(s[:])
	h.Set(sseCustomerAlgorithm, "AES256")
	h.Set(sseCustomerKey, base64.StdEncoding.EncodeToString(s[:]))
	h.Set(sseCustomerKeyMD5, base64.StdEncoding.EncodeToString(keyMD5[:]))
}

type ssecCopy [32]byte

func (s ssecCopy) Type() Type { return SSEC }

func (s ssecCopy) Marshal(h http.Header) {
	keyMD5 := md5.Sum(s[:])
	h.Set(sseCopyCustomerAlgorithm, "AES256")
	h.Set(sseCopyCustomerKey, base64.StdEncoding.EncodeToString(s[:]))
	h.Set(sseCopyCustomerKeyMD5, base64.StdEncoding.EncodeToString(keyMD5[:]))
}

type s3 struct{}

func (s s3) Type() Type { return S3 }

func (s s3) Marshal(h http.Header) { h.Set(sseGenericHeader, "AES256") }

type kms struct {
	key        string
	context    []byte
	hasContext bool
}

func (s kms) Type() Type { return KMS }

func (s kms) Marshal(h http.Header) {
	h.Set(sseGenericHeader, "aws:kms")
	if s.key != "" {
		h.Set(sseKmsKeyID, s.key)
	}
	if s.hasContext {
		h.Set(sseEncryptionContext, base64.StdEncoding.EncodeToString(s.context))
	}
}
<|endoftext|>


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

package cbft

import (
	"reflect"
	"testing"
)

func TestVersionGTE(t *testing.T) {
	tests := []struct {
		x        string
		y        string
		expected bool
	}{
		{"0.0.0", "0.0.0", true},
		{"0.0.1", "0.0.0", true},
		{"3.0.1", "2.0", true},
		{"3.0.0", "3.0", true},
		{"2.0.0", "2.0", true},
		{"2.0.1", "2.0", true},
		{"2.0.0", "2.5", false},
		{"1.0", "1.0.0", false},
		{"0.0", "0.0.0", false},
		{"", "", false},
		{"0", "", false},
		{"0.0", "", false},
		{"", "0", false},
		{"", "0.0", false},
		{"hello", "hello", false},
		{"0", "hello", false},
		{"0.0", "hello", false},
		{"hello", "0", false},
		{"hello", "0.0", false},
	}

	for i, test := range tests {
		actual := VersionGTE(test.x, test.y)
		if actual != test.expected {
			t.Errorf("test: %d, expected: %v, when %s >= %s, got: %v",
				i, test.expected, test.x, test.y, actual)
		}
	}
}

func TestNewUUID(t *testing.T) {
	u0 := NewUUID()
	u1 := NewUUID()
	if u0 == "" || u1 == "" || u0 == u1 {
		t.Errorf("NewUUID() failed, %s, %s", u0, u1)
	}
}

func TestExponentialBackoffLoop(t *testing.T) {
	called := 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		return -1
	}, 0, 0, 0)
	if called != 1 {
		t.Errorf("expected 1 call")
	}

	called = 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		if called <= 1 {
			return 1
		}
		return -1
	}, 0, 0, 0)
	if called != 2 {
		t.Errorf("expected 2 calls")
	}

	called = 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		if called == 1 {
			return 1
		}
		if called == 2 {
			return 0
		}
		return -1
	}, 0, 0, 0)
	if called != 3 {
		t.Errorf("expected 2 calls")
	}

	called = 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		if called == 1 {
			return 1
		}
		if called == 2 {
			return 0
		}
		return -1
	}, 1, 100000.0, 1)
	if called != 3 {
		t.Errorf("expected 2 calls")
	}
}

func TestTimeoutCancelChan(t *testing.T) {
	c := TimeoutCancelChan(0)
	if c != nil {
		t.Errorf("expected nil")
	}
	c = TimeoutCancelChan(-1)
	if c != nil {
		t.Errorf("expected nil")
	}
	c = TimeoutCancelChan(1)
	if c == nil {
		t.Errorf("expected non-nil")
	}
	msg, ok := <-c
	if ok {
		t.Errorf("expected closed")
	}
	if msg != false {
		t.Errorf("expected false msg on closed")
	}
}

func TestTime(t *testing.T) {
	x := uint64(100)
	Time(func() error { return nil }, &x)
	if x < 100 {
		t.Errorf("expected x to be >= 100")
	}
}

func TestAtomicCopyMetrics(t *testing.T) {
	src := &DestStats{
		TotError:     1,
		TotRollback:  2,
		TimeRollback: 3,
	}
	dst := DestStats{}
	AtomicCopyMetrics(src, &dst, nil)
	if !reflect.DeepEqual(src, &dst) {
		t.Errorf("expected src == dst")
	}
	if dst.TotError != 1 ||
		dst.TotRollback != 2 ||
		dst.TimeRollback != 3 {
		t.Errorf("expected src == dst")
	}
}
<commit_msg>TestMustEncode()<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package cbft

import (
	"bytes"
	"reflect"
	"testing"
)

func TestVersionGTE(t *testing.T) {
	tests := []struct {
		x        string
		y        string
		expected bool
	}{
		{"0.0.0", "0.0.0", true},
		{"0.0.1", "0.0.0", true},
		{"3.0.1", "2.0", true},
		{"3.0.0", "3.0", true},
		{"2.0.0", "2.0", true},
		{"2.0.1", "2.0", true},
		{"2.0.0", "2.5", false},
		{"1.0", "1.0.0", false},
		{"0.0", "0.0.0", false},
		{"", "", false},
		{"0", "", false},
		{"0.0", "", false},
		{"", "0", false},
		{"", "0.0", false},
		{"hello", "hello", false},
		{"0", "hello", false},
		{"0.0", "hello", false},
		{"hello", "0", false},
		{"hello", "0.0", false},
	}

	for i, test := range tests {
		actual := VersionGTE(test.x, test.y)
		if actual != test.expected {
			t.Errorf("test: %d, expected: %v, when %s >= %s, got: %v",
				i, test.expected, test.x, test.y, actual)
		}
	}
}

func TestNewUUID(t *testing.T) {
	u0 := NewUUID()
	u1 := NewUUID()
	if u0 == "" || u1 == "" || u0 == u1 {
		t.Errorf("NewUUID() failed, %s, %s", u0, u1)
	}
}

func TestExponentialBackoffLoop(t *testing.T) {
	called := 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		return -1
	}, 0, 0, 0)
	if called != 1 {
		t.Errorf("expected 1 call")
	}

	called = 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		if called <= 1 {
			return 1
		}
		return -1
	}, 0, 0, 0)
	if called != 2 {
		t.Errorf("expected 2 calls")
	}

	called = 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		if called == 1 {
			return 1
		}
		if called == 2 {
			return 0
		}
		return -1
	}, 0, 0, 0)
	if called != 3 {
		t.Errorf("expected 2 calls")
	}

	called = 0
	ExponentialBackoffLoop("test", func() int {
		called += 1
		if called == 1 {
			return 1
		}
		if called == 2 {
			return 0
		}
		return -1
	}, 1, 100000.0, 1)
	if called != 3 {
		t.Errorf("expected 2 calls")
	}
}

func TestTimeoutCancelChan(t *testing.T) {
	c := TimeoutCancelChan(0)
	if c != nil {
		t.Errorf("expected nil")
	}
	c = TimeoutCancelChan(-1)
	if c != nil {
		t.Errorf("expected nil")
	}
	c = TimeoutCancelChan(1)
	if c == nil {
		t.Errorf("expected non-nil")
	}
	msg, ok := <-c
	if ok {
		t.Errorf("expected closed")
	}
	if msg != false {
		t.Errorf("expected false msg on closed")
	}
}

func TestTime(t *testing.T) {
	x := uint64(100)
	Time(func() error { return nil }, &x)
	if x < 100 {
		t.Errorf("expected x to be >= 100")
	}
}

func TestAtomicCopyMetrics(t *testing.T) {
	src := &DestStats{
		TotError:     1,
		TotRollback:  2,
		TimeRollback: 3,
	}
	dst := DestStats{}
	AtomicCopyMetrics(src, &dst, nil)
	if !reflect.DeepEqual(src, &dst) {
		t.Errorf("expected src == dst")
	}
	if dst.TotError != 1 ||
		dst.TotRollback != 2 ||
		dst.TimeRollback != 3 {
		t.Errorf("expected src == dst")
	}
}

func TestMustEncode(t *testing.T) {
	defer func() {
		r := recover()
		if r == nil {
			t.Errorf("expected must encode panic and recover")
		}
	}()
	mustEncode(&bytes.Buffer{}, func() {})
	t.Errorf("expected must encode panic")
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package hostdisk

import (
	"fmt"
	"os"
	"path"
	"syscall"

	"kubevirt.io/kubevirt/pkg/log"

	k8sv1 "k8s.io/api/core/v1"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/util/types"
)

const (
	pvcBaseDir                  = "/var/run/kubevirt-private/vmi-disks"
	EventReasonToleratedSmallPV = "ToleratedSmallPV"
	EventTypeToleratedSmallPV   = k8sv1.EventTypeNormal
)

func ReplacePVCByHostDisk(vmi *v1.VirtualMachineInstance, clientset kubecli.KubevirtClient) error {
	// If PVC is defined and it's not a BlockMode PVC, then it is replaced by HostDisk
	// Filesystem PersistenVolumeClaim is mounted into pod as directory from node filesystem
	for i := range vmi.Spec.Volumes {
		if volumeSource := &vmi.Spec.Volumes[i].VolumeSource; volumeSource.PersistentVolumeClaim != nil {

			pvc, exists, isBlockVolumePVC, err := types.IsPVCBlockFromClient(clientset, vmi.Namespace, volumeSource.PersistentVolumeClaim.ClaimName)
			if err != nil {
				return err
			} else if !exists {
				return fmt.Errorf("persistentvolumeclaim %v not found", volumeSource.PersistentVolumeClaim.ClaimName)
			} else if isBlockVolumePVC {
				continue
			}

			volumeSource.HostDisk = &v1.HostDisk{
				Path:     getPVCDiskImgPath(vmi.Spec.Volumes[i].Name),
				Type:     v1.HostDiskExistsOrCreate,
				Capacity: pvc.Status.Capacity[k8sv1.ResourceStorage],
			}
			// PersistenVolumeClaim is replaced by HostDisk
			volumeSource.PersistentVolumeClaim = nil
		}
	}
	return nil
}

func dirBytesAvailable(path string) (uint64, error) {
	var stat syscall.Statfs_t
	err := syscall.Statfs(path, &stat)
	if err != nil {
		return 0, err
	}
	return stat.Bavail * uint64(stat.Bsize), nil
}

func createSparseRaw(fullPath string, size int64) error {
	offset := size - 1
	f, _ := os.Create(fullPath)
	defer f.Close()
	_, err := f.WriteAt([]byte{0}, offset)
	if err != nil {
		return err
	}
	return nil
}

func getPVCDiskImgPath(volumeName string) string {
	return path.Join(pvcBaseDir, volumeName, "disk.img")
}

type DiskImgCreator struct {
	dirBytesAvailableFunc  func(path string) (uint64, error)
	notifier               k8sNotifier
	lessPVCSpaceToleration int
}

type k8sNotifier interface {
	SendK8sEvent(vmi *v1.VirtualMachineInstance, severity string, reason string, message string) error
}

func NewHostDiskCreator(notifier k8sNotifier, lessPVCSpaceToleration int) DiskImgCreator {
	return DiskImgCreator{
		dirBytesAvailableFunc:  dirBytesAvailable,
		notifier:               notifier,
		lessPVCSpaceToleration: lessPVCSpaceToleration,
	}
}

func (hdc *DiskImgCreator) setlessPVCSpaceToleration(toleration int) {
	hdc.lessPVCSpaceToleration = toleration
}

func (hdc DiskImgCreator) Create(vmi *v1.VirtualMachineInstance) error {
	for _, volume := range vmi.Spec.Volumes {
		if hostDisk := volume.VolumeSource.HostDisk; hostDisk != nil && hostDisk.Type == v1.HostDiskExistsOrCreate && hostDisk.Path != "" {
			if _, err := os.Stat(hostDisk.Path); os.IsNotExist(err) {
				availableSize, err := hdc.dirBytesAvailableFunc(path.Dir(hostDisk.Path))
				if err != nil {
					return err
				}
				requestedSize, _ := hostDisk.Capacity.AsInt64()
				diskSize := requestedSize
				if uint64(requestedSize) > availableSize {
					// Some storage provisioners provision less space than requested, due to filesystem overhead etc.
					// We tolerate some difference in requested and available capacity up to some degree.
					// This can be configured with the "pvc-tolerate-less-space-up-to-percent" parameter in the kubevirt-config ConfigMap.
					// It is provided as argument to virt-launcher.
					toleratedSize := requestedSize * (100 - int64(hdc.lessPVCSpaceToleration)) / 100
					if uint64(toleratedSize) > availableSize {
						return fmt.Errorf("unable to create %s, not enough space, demanded size %d B is bigger than available space %d B, also after taking %v %% toleration into account",
							hostDisk.Path, uint64(requestedSize), availableSize, hdc.lessPVCSpaceToleration)
					}
					diskSize = int64(availableSize)

					msg := fmt.Sprintf("PV size too small: expected %v B, found %v B. Using it anyway, it is within %v %% toleration", requestedSize, availableSize, hdc.lessPVCSpaceToleration)
					log.Log.Info(msg)
					err = hdc.notifier.SendK8sEvent(vmi, EventTypeToleratedSmallPV, EventReasonToleratedSmallPV, msg)
					if err != nil {
						log.Log.Reason(err).Warningf("Couldn't send k8s event for tolerated PV size: %v", err)
					}
				}
				err = createSparseRaw(hostDisk.Path, int64(diskSize))
				if err != nil {
					return err
				}
			} else if err != nil {
				return err
			}

		}
	}
	return nil
}
<commit_msg>when replacing pvc with hostdisk indicate if pvc was shared<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package hostdisk

import (
	"fmt"
	"os"
	"path"
	"syscall"

	"kubevirt.io/kubevirt/pkg/log"

	k8sv1 "k8s.io/api/core/v1"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/util/types"
)

const (
	pvcBaseDir                  = "/var/run/kubevirt-private/vmi-disks"
	EventReasonToleratedSmallPV = "ToleratedSmallPV"
	EventTypeToleratedSmallPV   = k8sv1.EventTypeNormal
)

func ReplacePVCByHostDisk(vmi *v1.VirtualMachineInstance, clientset kubecli.KubevirtClient) error {
	// If PVC is defined and it's not a BlockMode PVC, then it is replaced by HostDisk
	// Filesystem PersistenVolumeClaim is mounted into pod as directory from node filesystem
	for i := range vmi.Spec.Volumes {
		if volumeSource := &vmi.Spec.Volumes[i].VolumeSource; volumeSource.PersistentVolumeClaim != nil {

			pvc, exists, isBlockVolumePVC, err := types.IsPVCBlockFromClient(clientset, vmi.Namespace, volumeSource.PersistentVolumeClaim.ClaimName)
			if err != nil {
				return err
			} else if !exists {
				return fmt.Errorf("persistentvolumeclaim %v not found", volumeSource.PersistentVolumeClaim.ClaimName)
			} else if isBlockVolumePVC {
				continue
			}
			isSharedPvc := false
			for _, accessMode := range pvc.Spec.AccessModes {
				if accessMode == k8sv1.ReadWriteMany {
					isSharedPvc = true
					break
				}
			}

			volumeSource.HostDisk = &v1.HostDisk{
				Path:     getPVCDiskImgPath(vmi.Spec.Volumes[i].Name),
				Type:     v1.HostDiskExistsOrCreate,
				Capacity: pvc.Status.Capacity[k8sv1.ResourceStorage],
				Shared:   &isSharedPvc,
			}
			// PersistenVolumeClaim is replaced by HostDisk
			volumeSource.PersistentVolumeClaim = nil
		}
	}
	return nil
}

func dirBytesAvailable(path string) (uint64, error) {
	var stat syscall.Statfs_t
	err := syscall.Statfs(path, &stat)
	if err != nil {
		return 0, err
	}
	return stat.Bavail * uint64(stat.Bsize), nil
}

func createSparseRaw(fullPath string, size int64) error {
	offset := size - 1
	f, _ := os.Create(fullPath)
	defer f.Close()
	_, err := f.WriteAt([]byte{0}, offset)
	if err != nil {
		return err
	}
	return nil
}

func getPVCDiskImgPath(volumeName string) string {
	return path.Join(pvcBaseDir, volumeName, "disk.img")
}

type DiskImgCreator struct {
	dirBytesAvailableFunc  func(path string) (uint64, error)
	notifier               k8sNotifier
	lessPVCSpaceToleration int
}

type k8sNotifier interface {
	SendK8sEvent(vmi *v1.VirtualMachineInstance, severity string, reason string, message string) error
}

func NewHostDiskCreator(notifier k8sNotifier, lessPVCSpaceToleration int) DiskImgCreator {
	return DiskImgCreator{
		dirBytesAvailableFunc:  dirBytesAvailable,
		notifier:               notifier,
		lessPVCSpaceToleration: lessPVCSpaceToleration,
	}
}

func (hdc *DiskImgCreator) setlessPVCSpaceToleration(toleration int) {
	hdc.lessPVCSpaceToleration = toleration
}

func (hdc DiskImgCreator) Create(vmi *v1.VirtualMachineInstance) error {
	for _, volume := range vmi.Spec.Volumes {
		if hostDisk := volume.VolumeSource.HostDisk; hostDisk != nil && hostDisk.Type == v1.HostDiskExistsOrCreate && hostDisk.Path != "" {
			if _, err := os.Stat(hostDisk.Path); os.IsNotExist(err) {
				availableSize, err := hdc.dirBytesAvailableFunc(path.Dir(hostDisk.Path))
				if err != nil {
					return err
				}
				requestedSize, _ := hostDisk.Capacity.AsInt64()
				diskSize := requestedSize
				if uint64(requestedSize) > availableSize {
					// Some storage provisioners provision less space than requested, due to filesystem overhead etc.
					// We tolerate some difference in requested and available capacity up to some degree.
					// This can be configured with the "pvc-tolerate-less-space-up-to-percent" parameter in the kubevirt-config ConfigMap.
					// It is provided as argument to virt-launcher.
					toleratedSize := requestedSize * (100 - int64(hdc.lessPVCSpaceToleration)) / 100
					if uint64(toleratedSize) > availableSize {
						return fmt.Errorf("unable to create %s, not enough space, demanded size %d B is bigger than available space %d B, also after taking %v %% toleration into account",
							hostDisk.Path, uint64(requestedSize), availableSize, hdc.lessPVCSpaceToleration)
					}
					diskSize = int64(availableSize)

					msg := fmt.Sprintf("PV size too small: expected %v B, found %v B. Using it anyway, it is within %v %% toleration", requestedSize, availableSize, hdc.lessPVCSpaceToleration)
					log.Log.Info(msg)
					err = hdc.notifier.SendK8sEvent(vmi, EventTypeToleratedSmallPV, EventReasonToleratedSmallPV, msg)
					if err != nil {
						log.Log.Reason(err).Warningf("Couldn't send k8s event for tolerated PV size: %v", err)
					}
				}
				err = createSparseRaw(hostDisk.Path, int64(diskSize))
				if err != nil {
					return err
				}
			} else if err != nil {
				return err
			}

		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2011 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package search_test

import (
	. "camlistore.org/pkg/search"

	"bytes"
	"encoding/json"
	"testing"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/test"
)

type describeTest struct {
	setup func(fi *test.FakeIndex)

	blob  string // blobref to describe
	depth int

	expect map[string]interface{}
}

var owner = blobref.MustParse("abcown-123")

var describeTests = []describeTest{
	{
		func(fi *test.FakeIndex) {},
		"abc-555",
		1,
		map[string]interface{}{},
	},

	{
		func(fi *test.FakeIndex) {
			fi.AddMeta(blobref.MustParse("abc-555"), "image/jpeg", 999)
		},
		"abc-555",
		1,
		map[string]interface{}{
			"abc-555": map[string]interface{}{
				"blobRef":  "abc-555",
				"mimeType": "image/jpeg",
				"size":     999,
			},
		},
	},

	{
		func(fi *test.FakeIndex) {
			pn := blobref.MustParse("perma-123")
			fi.AddMeta(pn, "application/json; camliType=permanode", 123)
			fi.AddClaim(owner, pn, "set-attribute", "camliContent", "foo-232")
			fi.AddMeta(blobref.MustParse("foo-232"), "foo/bar", 878)

			// Test deleting all attributes
			fi.AddClaim(owner, pn, "add-attribute", "wont-be-present", "x")
			fi.AddClaim(owner, pn, "add-attribute", "wont-be-present", "y")
			fi.AddClaim(owner, pn, "del-attribute", "wont-be-present", "")

			// Test deleting a specific attribute.
			fi.AddClaim(owner, pn, "add-attribute", "only-delete-b", "a")
			fi.AddClaim(owner, pn, "add-attribute", "only-delete-b", "b")
			fi.AddClaim(owner, pn, "add-attribute", "only-delete-b", "c")
			fi.AddClaim(owner, pn, "del-attribute", "only-delete-b", "b")
		},
		"perma-123",
		2,
		map[string]interface{}{
			"foo-232": map[string]interface{}{
				"blobRef":  "foo-232",
				"mimeType": "foo/bar",
				"size":     878,
			},
			"perma-123": map[string]interface{}{
				"blobRef":   "perma-123",
				"mimeType":  "application/json; camliType=permanode",
				"camliType": "permanode",
				"size":      123,
				"permanode": map[string]interface{}{
					"attr": map[string]interface{}{
						"camliContent":  []string{"foo-232"},
						"only-delete-b": []string{"a", "c"},
					},
				},
			},
		},
	},
}

func TestDescribe(t *testing.T) {
	for testn, tt := range describeTests {
		idx := test.NewFakeIndex()
		tt.setup(idx)

		h := NewHandler(idx, owner)
		js := make(map[string]interface{})
		dr := h.NewDescribeRequest()
		dr.Describe(blobref.MustParse(tt.blob), tt.depth)
		dr.PopulateJSON(js)
		got, _ := json.MarshalIndent(js, "", "  ")
		want, _ := json.MarshalIndent(tt.expect, "", "  ")
		if !bytes.Equal(got, want) {
			t.Errorf("test %d:\nwant: %s\n got: %s", testn, string(want), string(got))
		}
	}
}
<commit_msg>Improve search handler tests.<commit_after>/*
Copyright 2011 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package search_test

import (
	. "camlistore.org/pkg/search"

	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/test"
)

type handlerTest struct {
	// setup is responsible for populating the index before the
	// handler is invoked.
	//
	// A FakeIndex is constructed and provided to setup and is
	// generally then returned as the Index to use, but an
	// alternate Index may be returned instead, in which case the
	// FakeIndex is not used.
	setup func(fi *test.FakeIndex) Index

	query string // the HTTP path + optional query suffix after "camli/search/"

	want map[string]interface{}
}

var owner = blobref.MustParse("abcown-123")

var handlerTests = []handlerTest{
	{
		setup: func(fi *test.FakeIndex) Index { return fi },
		query: "describe?blobref=eabc-555",
		want:  map[string]interface{}{},
	},

	{
		setup: func(fi *test.FakeIndex) Index {
			fi.AddMeta(blobref.MustParse("abc-555"), "image/jpeg", 999)
			return fi
		},
		query: "describe?blobref=abc-555",
		want: map[string]interface{}{
			"abc-555": map[string]interface{}{
				"blobRef":  "abc-555",
				"mimeType": "image/jpeg",
				"size":     999,
			},
		},
	},

	{
	setup: func(fi *test.FakeIndex) Index {
			pn := blobref.MustParse("perma-123")
			fi.AddMeta(pn, "application/json; camliType=permanode", 123)
			fi.AddClaim(owner, pn, "set-attribute", "camliContent", "foo-232")
			fi.AddMeta(blobref.MustParse("foo-232"), "foo/bar", 878)

			// Test deleting all attributes
			fi.AddClaim(owner, pn, "add-attribute", "wont-be-present", "x")
			fi.AddClaim(owner, pn, "add-attribute", "wont-be-present", "y")
			fi.AddClaim(owner, pn, "del-attribute", "wont-be-present", "")

			// Test deleting a specific attribute.
			fi.AddClaim(owner, pn, "add-attribute", "only-delete-b", "a")
			fi.AddClaim(owner, pn, "add-attribute", "only-delete-b", "b")
			fi.AddClaim(owner, pn, "add-attribute", "only-delete-b", "c")
			fi.AddClaim(owner, pn, "del-attribute", "only-delete-b", "b")
			return fi
		},
		query: "describe?blobref=perma-123",
		want: map[string]interface{}{
			"foo-232": map[string]interface{}{
				"blobRef":  "foo-232",
				"mimeType": "foo/bar",
				"size":     878,
			},
			"perma-123": map[string]interface{}{
				"blobRef":   "perma-123",
				"mimeType":  "application/json; camliType=permanode",
				"camliType": "permanode",
				"size":      123,
				"permanode": map[string]interface{}{
					"attr": map[string]interface{}{
						"camliContent":  []string{"foo-232"},
						"only-delete-b": []string{"a", "c"},
					},
				},
			},
		},
	},
}

func TestHandler(t *testing.T) {
	for testn, tt := range handlerTests {
		fakeIndex := test.NewFakeIndex()
		idx := tt.setup(fakeIndex)
		h := NewHandler(idx, owner)

		req, err := http.NewRequest("GET", "/camli/search/"+tt.query, nil)
		if err != nil {
			t.Fatalf("%d: bad query: %v", testn, err)
		}
		req.Header.Set("X-PrefixHandler-PathSuffix", req.URL.Path[1:])

		rr := httptest.NewRecorder()
		rr.Body = new(bytes.Buffer)

		h.ServeHTTP(rr, req)

		got := rr.Body.Bytes()
		want, _ := json.MarshalIndent(tt.want, "", "  ")
		want = append(want, '\n')
		if !bytes.Equal(got, want) {
			t.Errorf("test %d:\nwant: %s\n got: %s", testn, want, got)
		}
	}
}
<|endoftext|>


<commit_before>package redis

import (
	"strconv"
	"sync"
	"time"
	"github.com/BluePecker/JwtAuth/pkg/storage"
	"github.com/go-redis/redis"
	"crypto/md5"
	"encoding/hex"
	"github.com/BluePecker/JwtAuth/pkg/storage/redis/uri"
	"reflect"
	"github.com/Sirupsen/logrus"
	"github.com/kataras/iris/core/errors"
)

type (
	Redis struct {
		create time.Time
		mu     sync.RWMutex
		engine Client
	}

	Client interface {
		Ping() *redis.StatusCmd

		Close() error

		Pipelined(fn func(redis.Pipeliner) error) ([]redis.Cmder, error)

		ZScore(key, field string) *redis.FloatCmd

		ZRem(key string, members ... interface{}) *redis.IntCmd

		ZRange(key string, start, stop int64) *redis.StringSliceCmd
	}
)

func inject(from, target reflect.Value) {
	indirect := reflect.Indirect(target.Elem())
	for index := 0; index < from.Elem().NumField(); index++ {
		name := from.Elem().Type().Field(index).Name
		f1 := from.Elem().FieldByName(name)
		f2 := indirect.FieldByName(name)
		if f2.IsValid() {
			if f1.Type() == f2.Type() && f2.CanSet() {
				f2.Set(f1)
			}
		}
	}
}

func (r *Redis) Initializer(opts string) error {
	generic, err := uri.Parser(opts)
	if err != nil {
		return err
	}
	switch reflect.ValueOf(generic).Interface().(type) {
	case *redis.ClusterOptions:
		options := &redis.ClusterOptions{}
		inject(reflect.ValueOf(generic), reflect.ValueOf(options))
		r.engine = redis.NewClusterClient(options)
		break
	case *redis.Options:
		options := &redis.Options{}
		inject(reflect.ValueOf(generic), reflect.ValueOf(options))
		r.engine = redis.NewClient(options)
		break
	}
	statusCmd := r.engine.Ping()
	if statusCmd.Err() != nil {
		logrus.Error(statusCmd.Err())
		defer r.engine.Close()
	}
	return statusCmd.Err()
}

func (r *Redis) HSet(key, field string, value interface{}, maxLen, expire int64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	_, err := r.engine.Pipelined(func(p redis.Pipeliner) error {
		exp := time.Duration(expire) * time.Second
		tmp := jwtMd5(key)
		val := jwtMd5(tmp + field)
		score := expire + int64(time.Now().Unix())
		if cmd := p.ZAdd(tmp, redis.Z{Score: float64(score), Member: val}); cmd.Err() != nil {
			return cmd.Err()
		} else {
			if cmd := p.Expire(tmp, exp); cmd.Err() != nil {
				return cmd.Err()
			}
			if cmd := p.Set(val, value, exp); cmd.Err() != nil {
				p.ZRem(tmp, redis.Z{Score: float64(score), Member: val})
				return cmd.Err()
			}
		}
		if cmd := p.ZCard(tmp); cmd.Err() != nil {
			return cmd.Err()
		} else {
			if cmd.Val() > maxLen {
				if cmd := p.ZRange(tmp, 0, cmd.Val()-maxLen); cmd.Err() != nil {
					return cmd.Err()
				} else {
					p.Del(cmd.Val()...)
				}
				cmd = p.ZRemRangeByRank(tmp, 0, cmd.Val()-maxLen)
				if cmd.Err() != nil {
					return cmd.Err()
				}
			}
		}
		return nil
	})
	return err
}

func (r *Redis) HGet(key, field string) (string, float64, error) {
	r.mu.RLock()
	defer r.mu.RLock()
	tmp := jwtMd5(key)
	return r.hGet(tmp, jwtMd5(tmp+field))
}

func (r *Redis) hGet(key, field string) (string, float64, error) {
	if cmd := r.engine.ZScore(key, field); cmd.Err() != nil {
		return "", -1, cmd.Err()
	} else if cmd.Val() < float64(time.Now().Unix()) {
		if cmd := r.engine.ZRem(key, field); cmd.Err() != nil {
			return "", -1, cmd.Err()
		} else {
			return "", -1, errors.New("key has been expired.")
		}
	} else {
		cmd, err := r.engine.Pipelined(func(p redis.Pipeliner) error {
			if cmd := p.Get(field); cmd.Err() != nil {
				return cmd.Err()
			}
			if cmd := p.TTL(field); cmd.Err() != nil {
				return cmd.Err()
			}
			return nil
		})
		if err == nil {
			if ttl, err := strconv.ParseFloat(cmd[1].String(), 64); err != nil {
				return "", -1, err
			} else {
				return cmd[0].String(), ttl, nil
			}
		}
		return "", -1, err
	}
}

func (r *Redis) HScan(key string, do func(token string, ttl float64)) error {
	r.mu.RLock()
	defer r.mu.RLock()
	tmp := jwtMd5(key)
	if cmd := r.engine.ZRange(tmp, 0, -1); cmd.Err() != nil {
		return cmd.Err()
	} else {
		for _, field := range cmd.Val() {
			singed, ttl, err := r.hGet(key, field)
			if err == nil {
				do(singed, ttl)
			}
		}
		return nil
	}
}

func (r *Redis) HRem(key string, field ... string) error {
	r.mu.Lock()
	defer r.mu.Lock()
	var v1 []interface{}
	var v2 []string
	tmp := jwtMd5(key)
	for _, v := range field {
		v1 = append(v1, jwtMd5(tmp+v))
		v2 = append(v2, jwtMd5(tmp+v))
	}
	_, err := r.engine.Pipelined(func(p redis.Pipeliner) error {
		p.ZRem(tmp, v1...)
		p.Del(v2...)
		return nil
	})
	return err
}

func jwtMd5(key string) string {
	hash := md5.New()
	hash.Write([]byte(key))
	return hex.EncodeToString(hash.Sum([]byte("jwt#")))
}

func init() {
	storage.Register("redis", &Redis{})
}
<commit_msg>fix bug<commit_after>package redis

import (
	"strconv"
	"sync"
	"time"
	"github.com/BluePecker/JwtAuth/pkg/storage"
	"github.com/go-redis/redis"
	"crypto/md5"
	"encoding/hex"
	"github.com/BluePecker/JwtAuth/pkg/storage/redis/uri"
	"reflect"
	"github.com/Sirupsen/logrus"
	"github.com/kataras/iris/core/errors"
)

type (
	Redis struct {
		create time.Time
		mu     sync.RWMutex
		engine Client
	}

	Client interface {
		Ping() *redis.StatusCmd

		Close() error

		Pipelined(fn func(redis.Pipeliner) error) ([]redis.Cmder, error)

		ZScore(key, field string) *redis.FloatCmd

		ZRem(key string, members ... interface{}) *redis.IntCmd

		ZRange(key string, start, stop int64) *redis.StringSliceCmd
	}
)

func inject(from, target reflect.Value) {
	indirect := reflect.Indirect(target.Elem())
	for index := 0; index < from.Elem().NumField(); index++ {
		name := from.Elem().Type().Field(index).Name
		f1 := from.Elem().FieldByName(name)
		f2 := indirect.FieldByName(name)
		if f2.IsValid() {
			if f1.Type() == f2.Type() && f2.CanSet() {
				f2.Set(f1)
			}
		}
	}
}

func (r *Redis) Initializer(opts string) error {
	generic, err := uri.Parser(opts)
	if err != nil {
		return err
	}
	switch reflect.ValueOf(generic).Interface().(type) {
	case *redis.ClusterOptions:
		options := &redis.ClusterOptions{}
		inject(reflect.ValueOf(generic), reflect.ValueOf(options))
		r.engine = redis.NewClusterClient(options)
		break
	case *redis.Options:
		options := &redis.Options{}
		inject(reflect.ValueOf(generic), reflect.ValueOf(options))
		r.engine = redis.NewClient(options)
		break
	}
	statusCmd := r.engine.Ping()
	if statusCmd.Err() != nil {
		logrus.Error(statusCmd.Err())
		defer r.engine.Close()
	}
	return statusCmd.Err()
}

func (r *Redis) HSet(key, field string, value interface{}, maxLen, expire int64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	_, err := r.engine.Pipelined(func(p redis.Pipeliner) error {
		exp := time.Duration(expire) * time.Second
		tmp := jwtMd5(key)
		val := jwtMd5(tmp + field)
		score := expire + int64(time.Now().Unix())
		if cmd := p.ZAdd(tmp, redis.Z{Score: float64(score), Member: val}); cmd.Err() != nil {
			return cmd.Err()
		} else {
			if cmd := p.Expire(tmp, exp); cmd.Err() != nil {
				return cmd.Err()
			}
			if cmd := p.Set(val, value, exp); cmd.Err() != nil {
				p.ZRem(tmp, redis.Z{Score: float64(score), Member: val})
				return cmd.Err()
			}
		}
		if cmd := p.ZCard(tmp); cmd.Err() != nil {
			return cmd.Err()
		} else {
			if cmd.Val() > maxLen {
				if cmd := p.ZRange(tmp, 0, cmd.Val()-maxLen); cmd.Err() != nil {
					return cmd.Err()
				} else {
					p.Del(cmd.Val()...)
				}
				cmd = p.ZRemRangeByRank(tmp, 0, cmd.Val()-maxLen)
				if cmd.Err() != nil {
					return cmd.Err()
				}
			}
		}
		return nil
	})
	return err
}

func (r *Redis) HGet(key, field string) (string, float64, error) {
	r.mu.RLock()
	defer r.mu.RLock()
	tmp := jwtMd5(key)
	return r.hGet(tmp, jwtMd5(tmp+field))
}

func (r *Redis) hGet(key, field string) (string, float64, error) {
	if cmd := r.engine.ZScore(key, field); cmd.Err() != nil {
		return "", -1, cmd.Err()
	} else if cmd.Val() < float64(time.Now().Unix()) {
		if cmd := r.engine.ZRem(key, field); cmd.Err() != nil {
			return "", -1, cmd.Err()
		} else {
			return "", -1, errors.New("key has been expired.")
		}
	} else {
		cmd, err := r.engine.Pipelined(func(p redis.Pipeliner) error {
			if cmd := p.Get(field); cmd.Err() != nil {
				return cmd.Err()
			}
			if cmd := p.TTL(field); cmd.Err() != nil {
				return cmd.Err()
			}
			return nil
		})
		if err == nil {
			if ttl, err := strconv.ParseFloat(cmd[1].String(), 64); err != nil {
				return "", -1, err
			} else {
				return cmd[0].String(), ttl, nil
			}
		}
		return "", -1, err
	}
}

func (r *Redis) HScan(key string, do func(token string, ttl float64)) error {
	r.mu.RLock()
	defer r.mu.RLock()
	tmp := jwtMd5(key)
	if cmd := r.engine.ZRange(tmp, 0, -1); cmd.Err() != nil {
		return cmd.Err()
	} else {
		for _, field := range cmd.Val() {
			singed, ttl, err := r.hGet(tmp, field)
			if err == nil {
				do(singed, ttl)
			}
		}
		return nil
	}
}

func (r *Redis) HRem(key string, field ... string) error {
	r.mu.Lock()
	defer r.mu.Lock()
	var v1 []interface{}
	var v2 []string
	tmp := jwtMd5(key)
	for _, v := range field {
		v1 = append(v1, jwtMd5(tmp+v))
		v2 = append(v2, jwtMd5(tmp+v))
	}
	_, err := r.engine.Pipelined(func(p redis.Pipeliner) error {
		p.ZRem(tmp, v1...)
		p.Del(v2...)
		return nil
	})
	return err
}

func jwtMd5(key string) string {
	hash := md5.New()
	hash.Write([]byte(key))
	return hex.EncodeToString(hash.Sum([]byte("jwt#")))
}

func init() {
	storage.Register("redis", &Redis{})
}
<|endoftext|>