<commit_before>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"errors"
	"fmt"
	"math"
	"testing"
	"time"
)

type stringer struct{}

func (s *stringer) String() string {
	return "data"
}

func TestEvent(t *testing.T) {
	defer resetEventTypes()

	now := time.Now()
	tStr := now.UTC().Format(TimeFormat)
	tStrNano := now.UTC().Format(time.RFC3339Nano)

	var msgTests = []struct {
		event        Event
		expected     string
		expectedJSON string
	}{
		{Event{DebugEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message6", 0},
			tStr + " [Debug] tag1, tag2, tag3: Message6, 0",
			`{"type": "Debug", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message6", "data": "0"}`},
		{Event{InfoEvent, now, Tags{"tag1", "tag2"}, "Message4", []byte("data")},
			tStr + " [Info] tag1, tag2: Message4, data",
			`{"type": "Info", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2"], ` +
				`"message": "Message4", "data": "data"}`},
		{Event{WarnEvent, now, Tags{"tag1"}, "Message3", &stringer{}},
			tStr + " [Warn] tag1: Message3, data",
			`{"type": "Warn", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message3", "data": "data"}`},
		{Event{ErrorEvent, now, Tags{"tag1"}, "Message2", "data"},
			tStr + " [Error] tag1: Message2, data",
			`{"type": "Error", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message2", "data": "data"}`},
		{Event{FatalEvent, now, Tags{}, "Message1", nil},
			tStr + " [Fatal] : Message1",
			`{"type": "Fatal", "timestamp": "` + tStrNano + `", "tags": [], ` +
				`"message": "Message1"}`},
		{Event{ThumbEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message5", errors.New("error data")},
			tStr + " [Thumb] tag1, tag2, tag3: Message5, error data",
			`{"type": "Thumb", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message5", "data": "error data"}`},
		{Event{NewEventType("My-event-type"), now, Tags{"tag1"}, "Message7", nil},
			tStr + " [My-event-type] tag1: Message7",
			`{"type": "My-event-type", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message7"}`},
		{Event{NewEventType(`my-"event"-type`), now, Tags{`tag"1"`}, "Message7", `"`},
			tStr + " [my-\"event\"-type] tag\"1\": Message7, \"",
			`{"type": "my-\"event\"-type", "timestamp": "` + tStrNano + `", "tags": ["tag\"1\""], ` +
				`"message": "Message7", "data": "\""}`},
	}

	for _, test := range msgTests {
		got, gotBytes := test.event.String(), string(test.event.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected Event(%v).Bytes() to return %q, but got %q",
				test.event, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected Event(%v).String() to return %q, but got %q",
				test.event, test.expected, got)
		}

		if json, err := test.event.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.event, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected Event(%v).MarshalJSON() to return %q, but got %q",
				test.event, test.expectedJSON, got)
		}
	}
}

func TestEventType(t *testing.T) {
	defer resetEventTypes()

	tests := []struct {
		eventType    EventType
		expected     string
		expectedJSON string
	}{
		{DebugEvent, "Debug", `"Debug"`},
		{ThumbEvent, "Thumb", `"Thumb"`},
		{InfoEvent, "Info", `"Info"`},
		{WarnEvent, "Warn", `"Warn"`},
		{ErrorEvent, "Error", `"Error"`},
		{FatalEvent, "Fatal", `"Fatal"`},
		{LogEvent, "Log", `"Log"`},
		{EventType(255), "EventType(255)", `"EventType(255)"`},
		{NewEventType("my-event-type"), "my-event-type", `"my-event-type"`},
		{NewEventType("my-\"event\"-type"), "my-\"event\"-type", `"my-\"event\"-type"`},
	}

	for _, test := range tests {
		got, gotBytes := test.eventType.String(), string(test.eventType.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected EventType(%v).Bytes() to return %q, but got %q",
				test.eventType, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected EventType(%v).String() to return %q, but got %q",
				test.eventType, test.expected, got)
		}

		if json, err := test.eventType.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.eventType, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected EventType(%v).MarshalJSON() to return %q, but got %q",
				test.eventType, test.expectedJSON, got)
		}

		var e = EventType(0)
		var gotEventType = &e

		if test.eventType == EventType(255) {
			expected := ErrEventTypeUnknown.Error()

			if err := gotEventType.UnmarshalText([]byte(test.expected)); err == nil {
				t.Fatalf("Expected the text unmarshalling of %v to return an error",
					test.eventType)
			} else if got := err.Error(); got != expected {
				t.Fatalf("Expected the returned error to be %s, but got %s",
					expected, got)
			}

			if err := gotEventType.UnmarshalJSON([]byte(test.expected)); err == nil {
				t.Fatalf("Expected the JSON unmarshalling of %v to return an error",
					test.eventType)
			} else if got := err.Error(); got != expected {
				t.Fatalf("Expected the returned error to be %s, but got %s",
					expected, got)
			}

			continue
		}

		if err := gotEventType.UnmarshalText([]byte(test.expected)); err != nil {
			t.Fatalf("Unexpected error text unmarshalling %s", err.Error())
		} else if *gotEventType != test.eventType {
			t.Fatalf("Expected the event type to be %v, but got %v",
				test.eventType, gotEventType)
		}

		if err := gotEventType.UnmarshalJSON([]byte(test.expectedJSON)); err != nil {
			t.Fatalf("Unexpected error JSON unmarshalling %s", err.Error())
		} else if *gotEventType != test.eventType {
			t.Fatalf("Expected the event type to be %v, but got %v",
				test.eventType, gotEventType)
		}
	}
}

func TestNewEventTypeWithEmptyName(t *testing.T) {
	defer resetEventTypes()

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal(`Expected NewEventType("") to panic, but it didn't`)
		}

		expected := "logger: EventType name can't be empty"
		got := recv.(string)
		if got != expected {
			t.Fatalf("Expected panic value to be %s, but got %s", expected, got)
		}
	}()

	NewEventType("")
}

func TestNewEventTypeNotUnique(t *testing.T) {
	defer resetEventTypes()

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal(`Expected a panic, but didn't get one`)
		}

		expected := "logger: EventType must be unique"
		got := recv.(string)
		if got != expected {
			t.Fatalf("Expected panic value to be %s, but got %s", expected, got)
		}
	}()

	NewEventType("my-event-type")
	NewEventType("my-event-type")
}

var (
	// Minus builtin event types.
	maxCostumEventTypes = math.MaxUint16 - len(eventTypeIndices)

	oldEventTypeNames   = eventTypeNames
	oldEventTypeIndices = eventTypeIndices
)

func resetEventTypes() {
	eventTypeNames = oldEventTypeNames
	eventTypeIndices = oldEventTypeIndices
}

func TestNewEventTypeLimit(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping TestNewEventTypeLimit in short mode")
	}
	defer resetEventTypes()

	for i := 1; i <= maxCostumEventTypes; i++ {
		expected := fmt.Sprintf("EventType-%d", i)
		eventType := NewEventType(expected)
		if got := eventType.String(); got != expected {
			t.Fatalf("Expected NewEventType(%q).String() to return %q, but got %q",
				expected, expected, got)
		}
	}

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a panic after creating 65528 log levels, but didn't get one")
		}

		got, ok := recv.(string)
		if !ok {
			t.Fatalf("Expected the recoverd panic to be a string, but it's %v", recv)
		}

		expected := "logger: can't have more then 65535 EventTypes"
		if got != expected {
			t.Fatalf("Expected the recoverd panic to be %s, but got %s", expected, got)
		}
	}()

	NewEventType(fmt.Sprintf("EventType-%d", maxCostumEventTypes+1))
}
<commit_msg>Split up TestEventType<commit_after>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"errors"
	"fmt"
	"math"
	"testing"
	"time"
)

type stringer struct{}

func (s *stringer) String() string {
	return "data"
}

func TestEvent(t *testing.T) {
	defer resetEventTypes()

	now := time.Now()
	tStr := now.UTC().Format(TimeFormat)
	tStrNano := now.UTC().Format(time.RFC3339Nano)

	var msgTests = []struct {
		event        Event
		expected     string
		expectedJSON string
	}{
		{Event{DebugEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message6", 0},
			tStr + " [Debug] tag1, tag2, tag3: Message6, 0",
			`{"type": "Debug", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message6", "data": "0"}`},
		{Event{InfoEvent, now, Tags{"tag1", "tag2"}, "Message4", []byte("data")},
			tStr + " [Info] tag1, tag2: Message4, data",
			`{"type": "Info", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2"], ` +
				`"message": "Message4", "data": "data"}`},
		{Event{WarnEvent, now, Tags{"tag1"}, "Message3", &stringer{}},
			tStr + " [Warn] tag1: Message3, data",
			`{"type": "Warn", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message3", "data": "data"}`},
		{Event{ErrorEvent, now, Tags{"tag1"}, "Message2", "data"},
			tStr + " [Error] tag1: Message2, data",
			`{"type": "Error", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message2", "data": "data"}`},
		{Event{FatalEvent, now, Tags{}, "Message1", nil},
			tStr + " [Fatal] : Message1",
			`{"type": "Fatal", "timestamp": "` + tStrNano + `", "tags": [], ` +
				`"message": "Message1"}`},
		{Event{ThumbEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message5", errors.New("error data")},
			tStr + " [Thumb] tag1, tag2, tag3: Message5, error data",
			`{"type": "Thumb", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message5", "data": "error data"}`},
		{Event{NewEventType("My-event-type"), now, Tags{"tag1"}, "Message7", nil},
			tStr + " [My-event-type] tag1: Message7",
			`{"type": "My-event-type", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message7"}`},
		{Event{NewEventType(`my-"event"-type`), now, Tags{`tag"1"`}, "Message7", `"`},
			tStr + " [my-\"event\"-type] tag\"1\": Message7, \"",
			`{"type": "my-\"event\"-type", "timestamp": "` + tStrNano + `", "tags": ["tag\"1\""], ` +
				`"message": "Message7", "data": "\""}`},
	}

	for _, test := range msgTests {
		got, gotBytes := test.event.String(), string(test.event.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected Event(%v).Bytes() to return %q, but got %q",
				test.event, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected Event(%v).String() to return %q, but got %q",
				test.event, test.expected, got)
		}

		if json, err := test.event.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.event, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected Event(%v).MarshalJSON() to return %q, but got %q",
				test.event, test.expectedJSON, got)
		}
	}
}

type eventTypeTest struct {
	EventType EventType
	Text      string
	JSON      string
}

func getEventTypesTests() []eventTypeTest {
	return []eventTypeTest{
		{DebugEvent, "Debug", `"Debug"`},
		{ThumbEvent, "Thumb", `"Thumb"`},
		{InfoEvent, "Info", `"Info"`},
		{WarnEvent, "Warn", `"Warn"`},
		{ErrorEvent, "Error", `"Error"`},
		{FatalEvent, "Fatal", `"Fatal"`},
		{LogEvent, "Log", `"Log"`},
		{EventType(255), "EventType(255)", `"EventType(255)"`},
		{NewEventType("my-event-type"), "my-event-type", `"my-event-type"`},
		{NewEventType("my-\"event\"-type"), "my-\"event\"-type", `"my-\"event\"-type"`},
	}
}

func TestEventTypeStringAndBytes(t *testing.T) {
	defer resetEventTypes()

	for _, test := range getEventTypesTests() {
		got, gotBytes := test.EventType.String(), string(test.EventType.Bytes())
		if gotBytes != test.Text {
			t.Errorf("Expected EventType(%v).Bytes() to return %q, but got %q",
				test.EventType, test.Text, gotBytes)
		} else if got != test.Text {
			t.Errorf("Expected EventType(%v).String() to return %q, but got %q",
				test.EventType, test.Text, got)
		}
	}
}

func TestEventTypeMarhslling(t *testing.T) {
	defer resetEventTypes()

	for _, test := range getEventTypesTests() {
		if json, err := test.EventType.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.EventType, err.Error())
		} else if got := string(json); got != test.JSON {
			t.Errorf("Expected EventType(%v).MarshalJSON() to return %q, but got %q",
				test.EventType, test.JSON, got)
		}
	}
}

func TestEventTypeUnmarshalling(t *testing.T) {
	defer resetEventTypes()

	for _, test := range getEventTypesTests() {
		var e = EventType(0)
		var gotEventType = &e

		var expectedError error = nil
		if _, ok := findEventType(test.EventType.String()); !ok {
			expectedError = ErrEventTypeUnknown
		}

		if err := gotEventType.UnmarshalText([]byte(test.Text)); err != expectedError {
			t.Fatalf("Expected EventType.UnmarshalText(%s) to return error %v, but got %v",
				test.Text, expectedError, err)
		} else if expectedError == nil && *gotEventType != test.EventType {
			t.Fatalf("Expected the event type to be %v, but got %v",
				test.EventType, gotEventType)
		}

		if err := gotEventType.UnmarshalJSON([]byte(test.JSON)); err != expectedError {
			t.Fatalf("Expected EventType.UnmarshalJSON(%s) to return error %v, but got %v",
				test.JSON, expectedError, err)
		} else if expectedError == nil && *gotEventType != test.EventType {
			t.Fatalf("Expected the event type to be %v, but got %v",
				test.EventType, gotEventType)
		}
	}
}

func TestNewEventTypeWithEmptyName(t *testing.T) {
	defer resetEventTypes()

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal(`Expected NewEventType("") to panic, but it didn't`)
		}

		expected := "logger: EventType name can't be empty"
		got := recv.(string)
		if got != expected {
			t.Fatalf("Expected panic value to be %s, but got %s", expected, got)
		}
	}()

	NewEventType("")
}

func TestNewEventTypeNotUnique(t *testing.T) {
	defer resetEventTypes()

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal(`Expected a panic, but didn't get one`)
		}

		expected := "logger: EventType must be unique"
		got := recv.(string)
		if got != expected {
			t.Fatalf("Expected panic value to be %s, but got %s", expected, got)
		}
	}()

	NewEventType("my-event-type")
	NewEventType("my-event-type")
}

var (
	// Minus builtin event types.
	maxCostumEventTypes = math.MaxUint16 - len(eventTypeIndices)

	oldEventTypeNames   = eventTypeNames
	oldEventTypeIndices = eventTypeIndices
)

func resetEventTypes() {
	eventTypeNames = oldEventTypeNames
	eventTypeIndices = oldEventTypeIndices
}

func TestNewEventTypeLimit(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping TestNewEventTypeLimit in short mode")
	}
	defer resetEventTypes()

	for i := 1; i <= maxCostumEventTypes; i++ {
		expected := fmt.Sprintf("EventType-%d", i)
		eventType := NewEventType(expected)
		if got := eventType.String(); got != expected {
			t.Fatalf("Expected NewEventType(%q).String() to return %q, but got %q",
				expected, expected, got)
		}
	}

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a panic after creating 65528 log levels, but didn't get one")
		}

		got, ok := recv.(string)
		if !ok {
			t.Fatalf("Expected the recoverd panic to be a string, but it's %v", recv)
		}

		expected := "logger: can't have more then 65535 EventTypes"
		if got != expected {
			t.Fatalf("Expected the recoverd panic to be %s, but got %s", expected, got)
		}
	}()

	NewEventType(fmt.Sprintf("EventType-%d", maxCostumEventTypes+1))
}
<|endoftext|>


<commit_before>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"io"
	"sync"
	"time"

	"google.golang.org/grpc"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
)

func init() {
	Commands = append(Commands, &commandEcEncode{})
}

type commandEcEncode struct {
}

func (c *commandEcEncode) Name() string {
	return "ec.encode"
}

func (c *commandEcEncode) Help() string {
	return `apply erasure coding to a volume

	ec.encode [-collection=""] [-fullPercent=95] [-quietFor=1h]
	ec.encode [-collection=""] [-volumeId=<volume_id>]

	This command will:
	1. freeze one volume
	2. apply erasure coding to the volume
	3. move the encoded shards to multiple volume servers

	The erasure coding is 10.4. So ideally you have more than 14 volume servers, and you can afford
	to lose 4 volume servers.

	If the number of volumes are not high, the worst case is that you only have 4 volume servers,
	and the shards are spread as 4,4,3,3, respectively. You can afford to lose one volume server.

	If you only have less than 4 volume servers, with erasure coding, at least you can afford to
	have 4 corrupted shard files.

`
}

func (c *commandEcEncode) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	encodeCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeId := encodeCommand.Int("volumeId", 0, "the volume id")
	collection := encodeCommand.String("collection", "", "the collection name")
	fullPercentage := encodeCommand.Float64("fullPercent", 95, "the volume reaches the percentage of max volume size")
	quietPeriod := encodeCommand.Duration("quietFor", time.Hour, "select volumes without no writes for this period")
	parallelCopy := encodeCommand.Bool("parallelCopy", true, "copy shards in parallel")
	if err = encodeCommand.Parse(args); err != nil {
		return nil
	}

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	vid := needle.VolumeId(*volumeId)

	// volumeId is provided
	if vid != 0 {
		return doEcEncode(commandEnv, *collection, vid, *parallelCopy)
	}

	// apply to all volumes in the collection
	volumeIds, err := collectVolumeIdsForEcEncode(commandEnv, *collection, *fullPercentage, *quietPeriod)
	if err != nil {
		return err
	}
	fmt.Printf("ec encode volumes: %v\n", volumeIds)
	for _, vid := range volumeIds {
		if err = doEcEncode(commandEnv, *collection, vid, *parallelCopy); err != nil {
			return err
		}
	}

	return nil
}

func doEcEncode(commandEnv *CommandEnv, collection string, vid needle.VolumeId, parallelCopy bool) (err error) {
	// find volume location
	locations, found := commandEnv.MasterClient.GetLocations(uint32(vid))
	if !found {
		return fmt.Errorf("volume %d not found", vid)
	}

	// fmt.Printf("found ec %d shards on %v\n", vid, locations)

	// mark the volume as readonly
	err = markVolumeReplicasWritable(commandEnv.option.GrpcDialOption, vid, locations, false)
	if err != nil {
		return fmt.Errorf("mark volume %d as readonly on %s: %v", vid, locations[0].Url, err)
	}

	// generate ec shards
	err = generateEcShards(commandEnv.option.GrpcDialOption, vid, collection, locations[0].ServerAddress())
	if err != nil {
		return fmt.Errorf("generate ec shards for volume %d on %s: %v", vid, locations[0].Url, err)
	}

	// balance the ec shards to current cluster
	err = spreadEcShards(commandEnv, vid, collection, locations, parallelCopy)
	if err != nil {
		return fmt.Errorf("spread ec shards for volume %d from %s: %v", vid, locations[0].Url, err)
	}

	return nil
}

func generateEcShards(grpcDialOption grpc.DialOption, volumeId needle.VolumeId, collection string, sourceVolumeServer pb.ServerAddress) error {

	fmt.Printf("generateEcShards %s %d on %s ...\n", collection, volumeId, sourceVolumeServer)

	err := operation.WithVolumeServerClient(sourceVolumeServer, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		_, genErr := volumeServerClient.VolumeEcShardsGenerate(context.Background(), &volume_server_pb.VolumeEcShardsGenerateRequest{
			VolumeId:   uint32(volumeId),
			Collection: collection,
		})
		return genErr
	})

	return err

}

func spreadEcShards(commandEnv *CommandEnv, volumeId needle.VolumeId, collection string, existingLocations []wdclient.Location, parallelCopy bool) (err error) {

	allEcNodes, totalFreeEcSlots, err := collectEcNodes(commandEnv, "")
	if err != nil {
		return err
	}

	if totalFreeEcSlots < erasure_coding.TotalShardsCount {
		return fmt.Errorf("not enough free ec shard slots. only %d left", totalFreeEcSlots)
	}
	allocatedDataNodes := allEcNodes
	if len(allocatedDataNodes) > erasure_coding.TotalShardsCount {
		allocatedDataNodes = allocatedDataNodes[:erasure_coding.TotalShardsCount]
	}

	// calculate how many shards to allocate for these servers
	allocatedEcIds := balancedEcDistribution(allocatedDataNodes)

	// ask the data nodes to copy from the source volume server
	copiedShardIds, err := parallelCopyEcShardsFromSource(commandEnv.option.GrpcDialOption, allocatedDataNodes, allocatedEcIds, volumeId, collection, existingLocations[0], parallelCopy)
	if err != nil {
		return err
	}

	// unmount the to be deleted shards
	err = unmountEcShards(commandEnv.option.GrpcDialOption, volumeId, existingLocations[0].ServerAddress(), copiedShardIds)
	if err != nil {
		return err
	}

	// ask the source volume server to clean up copied ec shards
	err = sourceServerDeleteEcShards(commandEnv.option.GrpcDialOption, collection, volumeId, existingLocations[0].ServerAddress(), copiedShardIds)
	if err != nil {
		return fmt.Errorf("source delete copied ecShards %s %d.%v: %v", existingLocations[0].Url, volumeId, copiedShardIds, err)
	}

	// ask the source volume server to delete the original volume
	for _, location := range existingLocations {
		fmt.Printf("delete volume %d from %s\n", volumeId, location.Url)
		err = deleteVolume(commandEnv.option.GrpcDialOption, volumeId, location.ServerAddress())
		if err != nil {
			return fmt.Errorf("deleteVolume %s volume %d: %v", location.Url, volumeId, err)
		}
	}

	return err

}

func parallelCopyEcShardsFromSource(grpcDialOption grpc.DialOption, targetServers []*EcNode, allocatedEcIds [][]uint32, volumeId needle.VolumeId, collection string, existingLocation wdclient.Location, parallelCopy bool) (actuallyCopied []uint32, err error) {

	fmt.Printf("parallelCopyEcShardsFromSource %d %s\n", volumeId, existingLocation.Url)

	var wg sync.WaitGroup
	shardIdChan := make(chan []uint32, len(targetServers))
	copyFunc := func(server *EcNode, allocatedEcShardIds []uint32) {
		defer wg.Done()
		copiedShardIds, copyErr := oneServerCopyAndMountEcShardsFromSource(grpcDialOption, server,
			allocatedEcShardIds, volumeId, collection, existingLocation.ServerAddress())
		if copyErr != nil {
			err = copyErr
		} else {
			shardIdChan <- copiedShardIds
			server.addEcVolumeShards(volumeId, collection, copiedShardIds)
		}
	}
	cleanupFunc := func(server *EcNode, allocatedEcShardIds []uint32) {
		if err := unmountEcShards(grpcDialOption, volumeId, pb.NewServerAddressFromDataNode(server.info), allocatedEcShardIds); err != nil {
			fmt.Printf("unmount aborted shards %d.%v on %s: %v\n", volumeId, allocatedEcShardIds, server.info.Id, err)
		}
		if err := sourceServerDeleteEcShards(grpcDialOption, collection, volumeId, pb.NewServerAddressFromDataNode(server.info), allocatedEcShardIds); err != nil {
			fmt.Printf("remove aborted shards %d.%v on %s: %v\n", volumeId, allocatedEcShardIds, server.info.Id, err)
		}
	}

	// maybe parallelize
	for i, server := range targetServers {
		if len(allocatedEcIds[i]) <= 0 {
			continue
		}

		wg.Add(1)
		if parallelCopy {
			go copyFunc(server, allocatedEcIds[i])
		} else {
			copyFunc(server, allocatedEcIds[i])
		}
	}
	wg.Wait()
	close(shardIdChan)

	if err != nil {
		for i, server := range targetServers {
			if len(allocatedEcIds[i]) <= 0 {
				continue
			}
			cleanupFunc(server, allocatedEcIds[i])
		}
		return nil, err
	}

	for shardIds := range shardIdChan {
		actuallyCopied = append(actuallyCopied, shardIds...)
	}

	return
}

func balancedEcDistribution(servers []*EcNode) (allocated [][]uint32) {
	allocated = make([][]uint32, len(servers))
	allocatedShardIdIndex := uint32(0)
	serverIndex := 0
	for allocatedShardIdIndex < erasure_coding.TotalShardsCount {
		if servers[serverIndex].freeEcSlot > 0 {
			allocated[serverIndex] = append(allocated[serverIndex], allocatedShardIdIndex)
			allocatedShardIdIndex++
		}
		serverIndex++
		if serverIndex >= len(servers) {
			serverIndex = 0
		}
	}

	return allocated
}

func collectVolumeIdsForEcEncode(commandEnv *CommandEnv, selectedCollection string, fullPercentage float64, quietPeriod time.Duration) (vids []needle.VolumeId, err error) {

	// collect topology information
	topologyInfo, volumeSizeLimitMb, err := collectTopologyInfo(commandEnv)
	if err != nil {
		return
	}

	quietSeconds := int64(quietPeriod / time.Second)
	nowUnixSeconds := time.Now().Unix()

	fmt.Printf("collect volumes quiet for: %d seconds\n", quietSeconds)

	vidMap := make(map[uint32]bool)
	eachDataNode(topologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		for _, diskInfo := range dn.DiskInfos {
			for _, v := range diskInfo.VolumeInfos {
				if v.Collection == selectedCollection && v.ModifiedAtSecond+quietSeconds < nowUnixSeconds {
					if float64(v.Size) > fullPercentage/100*float64(volumeSizeLimitMb)*1024*1024 {
						vidMap[v.Id] = true
					}
				}
			}
		}
	})

	for vid := range vidMap {
		vids = append(vids, needle.VolumeId(vid))
	}

	return
}
<commit_msg>adjust help message since both fullPercent and quietFor are needed.<commit_after>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"io"
	"sync"
	"time"

	"google.golang.org/grpc"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
)

func init() {
	Commands = append(Commands, &commandEcEncode{})
}

type commandEcEncode struct {
}

func (c *commandEcEncode) Name() string {
	return "ec.encode"
}

func (c *commandEcEncode) Help() string {
	return `apply erasure coding to a volume

	ec.encode [-collection=""] [-fullPercent=95 -quietFor=1h]
	ec.encode [-collection=""] [-volumeId=<volume_id>]

	This command will:
	1. freeze one volume
	2. apply erasure coding to the volume
	3. move the encoded shards to multiple volume servers

	The erasure coding is 10.4. So ideally you have more than 14 volume servers, and you can afford
	to lose 4 volume servers.

	If the number of volumes are not high, the worst case is that you only have 4 volume servers,
	and the shards are spread as 4,4,3,3, respectively. You can afford to lose one volume server.

	If you only have less than 4 volume servers, with erasure coding, at least you can afford to
	have 4 corrupted shard files.

`
}

func (c *commandEcEncode) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	encodeCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeId := encodeCommand.Int("volumeId", 0, "the volume id")
	collection := encodeCommand.String("collection", "", "the collection name")
	fullPercentage := encodeCommand.Float64("fullPercent", 95, "the volume reaches the percentage of max volume size")
	quietPeriod := encodeCommand.Duration("quietFor", time.Hour, "select volumes without no writes for this period")
	parallelCopy := encodeCommand.Bool("parallelCopy", true, "copy shards in parallel")
	if err = encodeCommand.Parse(args); err != nil {
		return nil
	}

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	vid := needle.VolumeId(*volumeId)

	// volumeId is provided
	if vid != 0 {
		return doEcEncode(commandEnv, *collection, vid, *parallelCopy)
	}

	// apply to all volumes in the collection
	volumeIds, err := collectVolumeIdsForEcEncode(commandEnv, *collection, *fullPercentage, *quietPeriod)
	if err != nil {
		return err
	}
	fmt.Printf("ec encode volumes: %v\n", volumeIds)
	for _, vid := range volumeIds {
		if err = doEcEncode(commandEnv, *collection, vid, *parallelCopy); err != nil {
			return err
		}
	}

	return nil
}

func doEcEncode(commandEnv *CommandEnv, collection string, vid needle.VolumeId, parallelCopy bool) (err error) {
	// find volume location
	locations, found := commandEnv.MasterClient.GetLocations(uint32(vid))
	if !found {
		return fmt.Errorf("volume %d not found", vid)
	}

	// fmt.Printf("found ec %d shards on %v\n", vid, locations)

	// mark the volume as readonly
	err = markVolumeReplicasWritable(commandEnv.option.GrpcDialOption, vid, locations, false)
	if err != nil {
		return fmt.Errorf("mark volume %d as readonly on %s: %v", vid, locations[0].Url, err)
	}

	// generate ec shards
	err = generateEcShards(commandEnv.option.GrpcDialOption, vid, collection, locations[0].ServerAddress())
	if err != nil {
		return fmt.Errorf("generate ec shards for volume %d on %s: %v", vid, locations[0].Url, err)
	}

	// balance the ec shards to current cluster
	err = spreadEcShards(commandEnv, vid, collection, locations, parallelCopy)
	if err != nil {
		return fmt.Errorf("spread ec shards for volume %d from %s: %v", vid, locations[0].Url, err)
	}

	return nil
}

func generateEcShards(grpcDialOption grpc.DialOption, volumeId needle.VolumeId, collection string, sourceVolumeServer pb.ServerAddress) error {

	fmt.Printf("generateEcShards %s %d on %s ...\n", collection, volumeId, sourceVolumeServer)

	err := operation.WithVolumeServerClient(sourceVolumeServer, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		_, genErr := volumeServerClient.VolumeEcShardsGenerate(context.Background(), &volume_server_pb.VolumeEcShardsGenerateRequest{
			VolumeId:   uint32(volumeId),
			Collection: collection,
		})
		return genErr
	})

	return err

}

func spreadEcShards(commandEnv *CommandEnv, volumeId needle.VolumeId, collection string, existingLocations []wdclient.Location, parallelCopy bool) (err error) {

	allEcNodes, totalFreeEcSlots, err := collectEcNodes(commandEnv, "")
	if err != nil {
		return err
	}

	if totalFreeEcSlots < erasure_coding.TotalShardsCount {
		return fmt.Errorf("not enough free ec shard slots. only %d left", totalFreeEcSlots)
	}
	allocatedDataNodes := allEcNodes
	if len(allocatedDataNodes) > erasure_coding.TotalShardsCount {
		allocatedDataNodes = allocatedDataNodes[:erasure_coding.TotalShardsCount]
	}

	// calculate how many shards to allocate for these servers
	allocatedEcIds := balancedEcDistribution(allocatedDataNodes)

	// ask the data nodes to copy from the source volume server
	copiedShardIds, err := parallelCopyEcShardsFromSource(commandEnv.option.GrpcDialOption, allocatedDataNodes, allocatedEcIds, volumeId, collection, existingLocations[0], parallelCopy)
	if err != nil {
		return err
	}

	// unmount the to be deleted shards
	err = unmountEcShards(commandEnv.option.GrpcDialOption, volumeId, existingLocations[0].ServerAddress(), copiedShardIds)
	if err != nil {
		return err
	}

	// ask the source volume server to clean up copied ec shards
	err = sourceServerDeleteEcShards(commandEnv.option.GrpcDialOption, collection, volumeId, existingLocations[0].ServerAddress(), copiedShardIds)
	if err != nil {
		return fmt.Errorf("source delete copied ecShards %s %d.%v: %v", existingLocations[0].Url, volumeId, copiedShardIds, err)
	}

	// ask the source volume server to delete the original volume
	for _, location := range existingLocations {
		fmt.Printf("delete volume %d from %s\n", volumeId, location.Url)
		err = deleteVolume(commandEnv.option.GrpcDialOption, volumeId, location.ServerAddress())
		if err != nil {
			return fmt.Errorf("deleteVolume %s volume %d: %v", location.Url, volumeId, err)
		}
	}

	return err

}

func parallelCopyEcShardsFromSource(grpcDialOption grpc.DialOption, targetServers []*EcNode, allocatedEcIds [][]uint32, volumeId needle.VolumeId, collection string, existingLocation wdclient.Location, parallelCopy bool) (actuallyCopied []uint32, err error) {

	fmt.Printf("parallelCopyEcShardsFromSource %d %s\n", volumeId, existingLocation.Url)

	var wg sync.WaitGroup
	shardIdChan := make(chan []uint32, len(targetServers))
	copyFunc := func(server *EcNode, allocatedEcShardIds []uint32) {
		defer wg.Done()
		copiedShardIds, copyErr := oneServerCopyAndMountEcShardsFromSource(grpcDialOption, server,
			allocatedEcShardIds, volumeId, collection, existingLocation.ServerAddress())
		if copyErr != nil {
			err = copyErr
		} else {
			shardIdChan <- copiedShardIds
			server.addEcVolumeShards(volumeId, collection, copiedShardIds)
		}
	}
	cleanupFunc := func(server *EcNode, allocatedEcShardIds []uint32) {
		if err := unmountEcShards(grpcDialOption, volumeId, pb.NewServerAddressFromDataNode(server.info), allocatedEcShardIds); err != nil {
			fmt.Printf("unmount aborted shards %d.%v on %s: %v\n", volumeId, allocatedEcShardIds, server.info.Id, err)
		}
		if err := sourceServerDeleteEcShards(grpcDialOption, collection, volumeId, pb.NewServerAddressFromDataNode(server.info), allocatedEcShardIds); err != nil {
			fmt.Printf("remove aborted shards %d.%v on %s: %v\n", volumeId, allocatedEcShardIds, server.info.Id, err)
		}
	}

	// maybe parallelize
	for i, server := range targetServers {
		if len(allocatedEcIds[i]) <= 0 {
			continue
		}

		wg.Add(1)
		if parallelCopy {
			go copyFunc(server, allocatedEcIds[i])
		} else {
			copyFunc(server, allocatedEcIds[i])
		}
	}
	wg.Wait()
	close(shardIdChan)

	if err != nil {
		for i, server := range targetServers {
			if len(allocatedEcIds[i]) <= 0 {
				continue
			}
			cleanupFunc(server, allocatedEcIds[i])
		}
		return nil, err
	}

	for shardIds := range shardIdChan {
		actuallyCopied = append(actuallyCopied, shardIds...)
	}

	return
}

func balancedEcDistribution(servers []*EcNode) (allocated [][]uint32) {
	allocated = make([][]uint32, len(servers))
	allocatedShardIdIndex := uint32(0)
	serverIndex := 0
	for allocatedShardIdIndex < erasure_coding.TotalShardsCount {
		if servers[serverIndex].freeEcSlot > 0 {
			allocated[serverIndex] = append(allocated[serverIndex], allocatedShardIdIndex)
			allocatedShardIdIndex++
		}
		serverIndex++
		if serverIndex >= len(servers) {
			serverIndex = 0
		}
	}

	return allocated
}

func collectVolumeIdsForEcEncode(commandEnv *CommandEnv, selectedCollection string, fullPercentage float64, quietPeriod time.Duration) (vids []needle.VolumeId, err error) {

	// collect topology information
	topologyInfo, volumeSizeLimitMb, err := collectTopologyInfo(commandEnv)
	if err != nil {
		return
	}

	quietSeconds := int64(quietPeriod / time.Second)
	nowUnixSeconds := time.Now().Unix()

	fmt.Printf("collect volumes quiet for: %d seconds\n", quietSeconds)

	vidMap := make(map[uint32]bool)
	eachDataNode(topologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		for _, diskInfo := range dn.DiskInfos {
			for _, v := range diskInfo.VolumeInfos {
				if v.Collection == selectedCollection && v.ModifiedAtSecond+quietSeconds < nowUnixSeconds {
					if float64(v.Size) > fullPercentage/100*float64(volumeSizeLimitMb)*1024*1024 {
						vidMap[v.Id] = true
					}
				}
			}
		}
	})

	for vid := range vidMap {
		vids = append(vids, needle.VolumeId(vid))
	}

	return
}
<|endoftext|>


<commit_before>// Copyright (c) Alex Ellis 2017. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package requests

// CreateFunctionRequest create a function in the swarm.
type CreateFunctionRequest struct {
	Service    string `json:"service"`
	Image      string `json:"image"`
	Network    string `json:"network"`
	EnvProcess string `json:"envProcess"`
}

type DeleteFunctionRequest struct {
	FunctionName string `json:"functionName"`
}

type AlexaSessionApplication struct {
	ApplicationId string `json:"applicationId"`
}

type AlexaSession struct {
	SessionId   string                  `json:"sessionId"`
	Application AlexaSessionApplication `json:"application"`
}

type AlexaIntent struct {
	Name string `json:"name"`
}

type AlexaRequest struct {
	Intent AlexaIntent `json:"intent"`
}

// AlexaRequestBody top-level request produced by Alexa SDK
type AlexaRequestBody struct {
	Session AlexaSession `json:"session"`
	Request AlexaRequest `json:"request"`
}

type PrometheusInnerAlertLabel struct {
	AlertName    string `json:"alertname"`
	FunctionName string `json:"function_name"`
}

type PrometheusInnerAlert struct {
	Status string                    `json:"status"`
	Labels PrometheusInnerAlertLabel `json:"labels"`
}

// PrometheusAlert as produced by AlertManager
type PrometheusAlert struct {
	Status   string                 `json:"status"`
	Receiver string                 `json:"receiver"`
	Alerts   []PrometheusInnerAlert `json:"alerts"`
}

// Function exported for system/functions endpoint
type Function struct {
	Name            string  `json:"name"`
	Image           string  `json:"image"`
	InvocationCount float64 `json:"invocationCount"`
	Replicas        uint64  `json:"replicas"`
}
<commit_msg>Add EnvVars<commit_after>// Copyright (c) Alex Ellis 2017. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package requests

// CreateFunctionRequest create a function in the swarm.
type CreateFunctionRequest struct {
	Service    string `json:"service"`
	Image      string `json:"image"`
	Network    string `json:"network"`

	// EnvProcess corresponds to the fprocess variable for your container watchdog.
	EnvProcess string `json:"envProcess"`

	// EnvVars provides overrides for functions.
	EnvVars map[string]string `json:"envVars"`
}

type DeleteFunctionRequest struct {
	FunctionName string `json:"functionName"`
}

type AlexaSessionApplication struct {
	ApplicationId string `json:"applicationId"`
}

type AlexaSession struct {
	SessionId   string                  `json:"sessionId"`
	Application AlexaSessionApplication `json:"application"`
}

type AlexaIntent struct {
	Name string `json:"name"`
}

type AlexaRequest struct {
	Intent AlexaIntent `json:"intent"`
}

// AlexaRequestBody top-level request produced by Alexa SDK
type AlexaRequestBody struct {
	Session AlexaSession `json:"session"`
	Request AlexaRequest `json:"request"`
}

type PrometheusInnerAlertLabel struct {
	AlertName    string `json:"alertname"`
	FunctionName string `json:"function_name"`
}

type PrometheusInnerAlert struct {
	Status string                    `json:"status"`
	Labels PrometheusInnerAlertLabel `json:"labels"`
}

// PrometheusAlert as produced by AlertManager
type PrometheusAlert struct {
	Status   string                 `json:"status"`
	Receiver string                 `json:"receiver"`
	Alerts   []PrometheusInnerAlert `json:"alerts"`
}

// Function exported for system/functions endpoint
type Function struct {
	Name            string  `json:"name"`
	Image           string  `json:"image"`
	InvocationCount float64 `json:"invocationCount"`
	Replicas        uint64  `json:"replicas"`
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"fmt"
	"log"
	"os"
	"sync"
	"unicode"
	"unicode/utf8"

	"github.com/uber/tchannel/golang"
	"github.com/uber/tchannel/golang/examples/keyvalue/gen-go/keyvalue"
	"github.com/uber/tchannel/golang/hyperbahn"
	"github.com/uber/tchannel/golang/thrift"
)

func main() {
	// Create a TChannel and listen for inbound connections.
	ch, err := tchannel.NewChannel("keyvalue", nil)
	if err != nil {
		log.Fatalf("Failed to create tchannel: %v", err)
	}
	ip, err := tchannel.ListenIP()
	if err != nil {
		log.Fatalf("Failed to find IP to Listen on: %v", err)
	}
	ch.ListenAndServe(ip.String() + ":12345")

	// Register both the KeyValue and Admin services.
	h := newKVHandler()
	server := thrift.NewServer(ch)
	server.Register(keyvalue.NewTChanKeyValueServer(h))
	server.Register(keyvalue.NewTChanAdminServer(h))

	config := hyperbahn.Configuration{InitialNodes: os.Args[1:]}
	if len(config.InitialNodes) == 0 {
		log.Fatalf("No Autobahn nodes to advertise with")
	}
	client, err := hyperbahn.NewClient(ch, config, nil)
	if err != nil {
		log.Fatalf("hyperbahn.NewClient failed: %v", err)
	}
	if err := client.Advertise(); err != nil {
		log.Fatalf("Hyperbahn advertise failed: %v", err)
	}

	// The service is now started up, run it till we receive a ctrl-c.
	log.Printf("KeyValue service has started on %v", ch.PeerInfo().HostPort)
	select {}
}

type kvHandler struct {
	mut  sync.RWMutex
	vals map[string]string
}

// NewKVHandler returns a new handler for the KeyValue service.
func newKVHandler() *kvHandler {
	return &kvHandler{vals: make(map[string]string)}
}

// Get returns the value stored for the given key.
func (h *kvHandler) Get(ctx thrift.Context, key string) (string, error) {
	if err := isValidKey(key); err != nil {
		return "", err
	}

	h.mut.RLock()
	defer h.mut.RUnlock()

	if val, ok := h.vals[key]; ok {
		return val, nil
	}

	return "", &keyvalue.KeyNotFound{Key: key}
}

// Set sets the value for a given key.
func (h *kvHandler) Set(ctx thrift.Context, key, value string) error {
	if err := isValidKey(key); err != nil {
		return err
	}

	h.mut.Lock()
	defer h.mut.Unlock()

	h.vals[key] = value
	// Example of how to use response headers. Normally, these values should be passed via result structs.
	ctx.SetResponseHeaders(map[string]string{"count": fmt.Sprint(len(h.vals))})
	return nil
}

// HealthCheck return the health status of this process.
func (h *kvHandler) HealthCheck(ctx thrift.Context) (string, error) {
	return "OK", nil
}

// ClearAll clears all the keys.
func (h *kvHandler) ClearAll(ctx thrift.Context) error {
	if !isAdmin(ctx) {
		return &keyvalue.NotAuthorized{}
	}

	h.mut.Lock()
	defer h.mut.Unlock()

	h.vals = make(map[string]string)
	return nil
}

func isValidKey(key string) error {
	r, _ := utf8.DecodeRuneInString(key)
	if !unicode.IsLetter(r) {
		return &keyvalue.InvalidKey{}
	}
	return nil
}

func isAdmin(ctx thrift.Context) bool {
	return ctx.Headers()["user"] == "admin"
}
<commit_msg>Improve keyvalue server example<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"fmt"
	"log"
	"os"
	"sync"
	"unicode"
	"unicode/utf8"

	"github.com/uber/tchannel/golang"
	"github.com/uber/tchannel/golang/examples/keyvalue/gen-go/keyvalue"
	"github.com/uber/tchannel/golang/hyperbahn"
	"github.com/uber/tchannel/golang/thrift"
)

func main() {
	// Create a TChannel and register the Thrift handlers.
	ch, err := tchannel.NewChannel("keyvalue", nil)
	if err != nil {
		log.Fatalf("Failed to create tchannel: %v", err)
	}

	// Register both the KeyValue and Admin services.
	// We can register multiple Thrift services on a single Hyperbahn service.
	h := newKVHandler()
	server := thrift.NewServer(ch)
	server.Register(keyvalue.NewTChanKeyValueServer(h))
	server.Register(keyvalue.NewTChanAdminServer(h))

	// Listen for connections on the external interface so we can receive connections.
	ip, err := tchannel.ListenIP()
	if err != nil {
		log.Fatalf("Failed to find IP to Listen on: %v", err)
	}
	// We use port 0 which asks the OS to assign any available port.
	// Static port allocations are not necessary for services on Hyperbahn.
	ch.ListenAndServe(fmt.Sprintf("%v:%v", ip, 0))

	// Advertising registers this service instance with Hyperbahn so
	// that Hyperbahn can route requests for "keyvalue" to us.
	config := hyperbahn.Configuration{InitialNodes: os.Args[1:]}
	if len(config.InitialNodes) == 0 {
		log.Fatalf("No Autobahn nodes to advertise with")
	}
	client, err := hyperbahn.NewClient(ch, config, nil)
	if err != nil {
		log.Fatalf("hyperbahn.NewClient failed: %v", err)
	}
	if err := client.Advertise(); err != nil {
		log.Fatalf("Hyperbahn advertise failed: %v", err)
	}

	// The service is now started up, run it till we receive a ctrl-c.
	log.Printf("KeyValue service has started on %v", ch.PeerInfo().HostPort)
	select {}
}

type kvHandler struct {
	mut  sync.RWMutex
	vals map[string]string
}

// NewKVHandler returns a new handler for the KeyValue service.
func newKVHandler() *kvHandler {
	return &kvHandler{vals: make(map[string]string)}
}

// Get returns the value stored for the given key.
func (h *kvHandler) Get(ctx thrift.Context, key string) (string, error) {
	if err := isValidKey(key); err != nil {
		return "", err
	}

	h.mut.RLock()
	defer h.mut.RUnlock()

	if val, ok := h.vals[key]; ok {
		return val, nil
	}

	return "", &keyvalue.KeyNotFound{Key: key}
}

// Set sets the value for a given key.
func (h *kvHandler) Set(ctx thrift.Context, key, value string) error {
	if err := isValidKey(key); err != nil {
		return err
	}

	h.mut.Lock()
	defer h.mut.Unlock()

	h.vals[key] = value
	// Example of how to use response headers. Normally, these values should be passed via result structs.
	ctx.SetResponseHeaders(map[string]string{"count": fmt.Sprint(len(h.vals))})
	return nil
}

// HealthCheck return the health status of this process.
func (h *kvHandler) HealthCheck(ctx thrift.Context) (string, error) {
	return "OK", nil
}

// ClearAll clears all the keys.
func (h *kvHandler) ClearAll(ctx thrift.Context) error {
	if !isAdmin(ctx) {
		return &keyvalue.NotAuthorized{}
	}

	h.mut.Lock()
	defer h.mut.Unlock()

	h.vals = make(map[string]string)
	return nil
}

func isValidKey(key string) error {
	r, _ := utf8.DecodeRuneInString(key)
	if !unicode.IsLetter(r) {
		return &keyvalue.InvalidKey{}
	}
	return nil
}

func isAdmin(ctx thrift.Context) bool {
	return ctx.Headers()["user"] == "admin"
}
<|endoftext|>


<commit_before>package youtube

import(
	"os/exec"
	"fmt"
	"strings"
)

func YoutubeDl(url string) (fileName string, err error) {
	format := selectVideoFormat(url)
	fmt.Println("format:", format)

	fmt.Println("youtube-dl", "-o", "./steamerDataDir/%(title)s-%(format)-%(id)s.%(ext)s", "-f", format, url)
	cmd := exec.Command("youtube-dl", "-o", "./steamerDataDir/%(title)s-%(format)s-%(id)s.%(ext)s", "-f", format, url)
	output, err := cmd.Output()

	if err != nil {
		return "", err
	}
	fmt.Println("download content: ", string(output))

	cmd = exec.Command("youtube-dl", "--get-filename", "-o", "./steamerDataDir/%(title)s-%(format)s-%(id)s.%(ext)s", "-f", format, url)
	output, err = cmd.Output()
	if err != nil {
		return "", err
	}

	fileName = strings.TrimSpace(string(output))
	fmt.Println("download fileName:", fileName)

	return fileName, nil
}

func selectVideoFormat(url string) string {
	fmt.Println("exec command youtube-dl -F", url)
	cmd := exec.Command("youtube-dl", "-F", url)
	output, err := cmd.Output();
	if( err != nil ){
		fmt.Println("get url formats error:", url)
	}
	fmt.Println("output contnent is:", string(output))
	lines := strings.Split(string(output), "\n")
	lastLine := lines[len(lines)-2]

	lastLineFields := strings.Fields(lastLine)
	return lastLineFields[0]

}<commit_msg>处理并返回视频信息<commit_after>package youtube

import (
	"encoding/json"
	"fmt"
	"os/exec"
	"strings"
)

type VideoQuality struct {
	FormatId   string
	Extension  string
	Resolution string
	Note       string
}

func YoutubeDl(url string) (fileName string, err error) {
	format := selectVideoFormat(url)
	fmt.Println("format:", format)

	fmt.Println("youtube-dl", "-o", "./steamerDataDir/%(title)s-%(format)-%(id)s.%(ext)s", "-f", format, url)
	cmd := exec.Command("youtube-dl", "-o", "./steamerDataDir/%(title)s-%(format)s-%(id)s.%(ext)s", "-f", format, url)
	output, err := cmd.Output()

	if err != nil {
		return "", err
	}
	fmt.Println("download content: ", string(output))

	cmd = exec.Command("youtube-dl", "--get-filename", "-o", "./steamerDataDir/%(title)s-%(format)s-%(id)s.%(ext)s", "-f", format, url)
	output, err = cmd.Output()
	if err != nil {
		return "", err
	}

	fileName = strings.TrimSpace(string(output))
	fmt.Println("download fileName:", fileName)

	return fileName, nil
}

func selectVideoFormat(url string) string {
	fmt.Println("exec command youtube-dl -F", url)
	cmd := exec.Command("youtube-dl", "-F", url)
	output, err := cmd.Output()
	if err != nil {
		fmt.Println("get url formats error:", url)
	}
	fmt.Println("output contnent is:", string(output))
	lines := strings.Split(string(output), "\n")
	lastLine := lines[len(lines)-2]

	lastLineFields := strings.Fields(lastLine)
	return lastLineFields[0]

}

func GetVideoQuality(url string) string {
	videoQualities := make([]VideoQuality, 0)
	//获取视频信息
	cmd := exec.Command("youtube-dl", "-F", url)
	output, err := cmd.Output()
	if err != nil {
		fmt.Println("获取视频信息失败")
	}
	//解析字符串信息
	outputArr := strings.Split(string(output), "\n")
	for i := 0; i < len(outputArr); i++ {

		conArr := strings.Fields(outputArr[i])

		if len(conArr) > 4 && conArr[3] == "DASH" && conArr[4] == "video" {

			conStr3 := ""
			for i := 3; i < len(conArr[3:])+3; i++ {
				conStr3 += conArr[i]
			}

			videoQualities = append(videoQualities, VideoQuality{
				FormatId:   conArr[0],
				Extension:  conArr[1],
				Resolution: conArr[2],
				Note:       conStr3,
			})
		}
	}
	str, _ := json.Marshal(videoQualities)
	return str
}
<|endoftext|>


<commit_before>package api

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/sessions"
	"io/ioutil"
	"net/http"
)

/*
//API Register
router := api.Default()
apiRouter := router.Group("/api")
apiRouter.RequirePermission(true)
apiStoragesRouter := apiRouter.Group("/storages")
apiStoragesRouter.RequirePermission(true)
apiStoragesRouter.PUT("/", func(c *api.Context, s *sessions.Session) {
	c.Abort(http.StatusOK)
})
apiStoragesRouter.GET("/", func(c *api.Context, s *sessions.Session) {
	c.Resolve(http.StatusOK, "123")
})
apiStoragesRouter.POST("/", func(c *api.Context, s *sessions.Session) {
	c.Resolve(http.StatusOK, nil)
})
apiStoragesRouter.DELETE("/", func(c *api.Context, s *sessions.Session) {
	c.Abort(http.StatusOK)
})
*/
type Engine struct {
	gin   *gin.Engine
	store *sessions.CookieStore
	addr  string
}

type EngineGroup struct {
	engine      *Engine
	children    []*EngineGroup
	routerGroup *gin.RouterGroup
	path        string
	handlerHash map[string][]localAPIHandler
	CORS        bool
}

func (this *EngineGroup) SetCORS(CORS bool) {
	this.CORS = CORS
}

func (this *EngineGroup) SetCORSAll(CORS bool) {
	this.SetCORS(CORS)
	for _, child := range this.children {
		child.SetCORSAll(CORS)
	}
}

type localAPIHandler func(*Context)

type APIHandler func(*Context, *sessions.Session)

func Default(name string, addr string) *EngineGroup {
	engine := &Engine{
		gin:   gin.Default(),
		store: sessions.NewCookieStore([]byte(fmt.Sprintf("%s-Server-Store", name))),
		addr:  addr,
	}
	router := &EngineGroup{
		engine:      engine,
		children:    make([]*EngineGroup, 0, 16),
		routerGroup: nil,
		path:        "",
		handlerHash: make(map[string][]localAPIHandler),
	}
	engine.gin.OPTIONS("/*", func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.AbortWithStatus(http.StatusOK)
	})
	engine.gin.NoMethod(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.AbortWithStatus(http.StatusMethodNotAllowed)
	})
	engine.gin.NotFound404(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.AbortWithStatus(http.StatusNotFound)
	})
	apiLocalSupportRegister(addr, router)
	return router
}

func (this *EngineGroup) Group(path string) *EngineGroup {
	var routerGroup *gin.RouterGroup
	if this.routerGroup != nil {
		routerGroup = this.routerGroup.Group(path)
	} else {
		routerGroup = this.engine.gin.Group(path)
	}
	child := &EngineGroup{
		engine:      this.engine,
		routerGroup: routerGroup,
		path:        path,
		handlerHash: make(map[string][]localAPIHandler),
	}
	this.children = append(this.children, child)
	return child
}

const (
	METHOD_PUT    = iota
	METHOD_GET    = iota
	METHOD_POST   = iota
	METHOD_DELETE = iota
	method_COUNT  = iota
)

func (this *EngineGroup) register(method int, path string, handler APIHandler) {
	apiList, has := this.handlerHash[path]
	if !has {
		apiList = make([]localAPIHandler, method_COUNT)
		this.handlerHash[path] = apiList
	}
	apiList[method] = func(c *Context) {
		session, _ := this.engine.store.Get(c.Request, API_SESSION_NAME)
		handler(c, session)
		session.Save(c.Request, c.Writer)
	}

	handlerImp := this.getHandlerImp(handler)
	switch method {
	case METHOD_PUT:
		if this.routerGroup != nil {
			this.routerGroup.PUT(path, handlerImp)
		} else {
			this.engine.gin.PUT(path, handlerImp)
		}
	case METHOD_GET:
		if this.routerGroup != nil {
			this.routerGroup.GET(path, handlerImp)
		} else {
			this.engine.gin.GET(path, handlerImp)
		}
	case METHOD_POST:
		if this.routerGroup != nil {
			this.routerGroup.POST(path, handlerImp)
		} else {
			this.engine.gin.POST(path, handlerImp)
		}
	case METHOD_DELETE:
		if this.routerGroup != nil {
			this.routerGroup.DELETE(path, handlerImp)
		} else {
			this.engine.gin.DELETE(path, handlerImp)
		}
	}
}

func (this *EngineGroup) PUT(path string, handler APIHandler) {
	this.register(METHOD_PUT, path, handler)
}

func (this *EngineGroup) GET(path string, handler APIHandler) {
	this.register(METHOD_GET, path, handler)
}

func (this *EngineGroup) POST(path string, handler APIHandler) {
	this.register(METHOD_POST, path, handler)
}

func (this *EngineGroup) DELETE(path string, handler APIHandler) {
	this.register(METHOD_DELETE, path, handler)
}

func (this *EngineGroup) getHandlerImp(handler APIHandler) gin.HandlerFunc {
	return func(c *gin.Context) {
		var data interface{}
		body, _ := ioutil.ReadAll(c.Request.Body)
		err := json.Unmarshal(body, &data)
		if err != nil {
			var buffer bytes.Buffer
			buffer.Write(body)
			c.Request.Body = ioutil.NopCloser(bufio.NewReader(&buffer))
		}
		context := &Context{
			code:    unresolvedCode,
			ret:     nil,
			Data:    data,
			Params:  c.Params,
			QParams: c.Request.URL.Query(),
			Request: c.Request,
			Writer:  c.Writer,
		}
		session, _ := this.engine.store.Get(c.Request, API_SESSION_NAME)
		handler(context, session)
		session.Save(c.Request, c.Writer)
		if this.CORS {
			c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
	        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
			c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		}
		switch context.code {
		case unresolvedCode:
			c.AbortWithStatus(http.StatusInternalServerError)
		case http.StatusMovedPermanently:
			c.Redirect(context.code, context.ret.(string))
		case http.StatusTemporaryRedirect:
			c.Redirect(context.code, context.ret.(string))
		default:
			switch context.ret.(type) {
			case nil:
				c.AbortWithStatus(context.code)
			case string:
				c.HTMLString(context.code, context.ret.(string))
				//c.String(context.code, context.ret.(string))
			default:
				c.JSON(context.code, context.ret)
			}
		}
	}
}

func (this *EngineGroup) Run() error {
	return this.engine.gin.Run(this.engine.addr)
}

func (this *EngineGroup) RunTLS(cert string, key string) error {
	return this.engine.gin.RunTLS(this.engine.addr, cert, key)
}
<commit_msg>bug fix<commit_after>package api

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/sessions"
	"io/ioutil"
	"net/http"
)

/*
//API Register
router := api.Default()
apiRouter := router.Group("/api")
apiRouter.RequirePermission(true)
apiStoragesRouter := apiRouter.Group("/storages")
apiStoragesRouter.RequirePermission(true)
apiStoragesRouter.PUT("/", func(c *api.Context, s *sessions.Session) {
	c.Abort(http.StatusOK)
})
apiStoragesRouter.GET("/", func(c *api.Context, s *sessions.Session) {
	c.Resolve(http.StatusOK, "123")
})
apiStoragesRouter.POST("/", func(c *api.Context, s *sessions.Session) {
	c.Resolve(http.StatusOK, nil)
})
apiStoragesRouter.DELETE("/", func(c *api.Context, s *sessions.Session) {
	c.Abort(http.StatusOK)
})
*/
type Engine struct {
	gin   *gin.Engine
	store *sessions.CookieStore
	addr  string
}

type EngineGroup struct {
	engine      *Engine
	children    []*EngineGroup
	routerGroup *gin.RouterGroup
	path        string
	handlerHash map[string][]localAPIHandler
	CORS        bool
}

func (this *EngineGroup) SetCORS(CORS bool) {
	this.CORS = CORS
}

func (this *EngineGroup) SetCORSAll(CORS bool) {
	this.SetCORS(CORS)
	for _, child := range this.children {
		child.SetCORSAll(CORS)
	}
}

type localAPIHandler func(*Context)

type APIHandler func(*Context, *sessions.Session)

func Default(name string, addr string) *EngineGroup {
	engine := &Engine{
		gin:   gin.Default(),
		store: sessions.NewCookieStore([]byte(fmt.Sprintf("%s-Server-Store", name))),
		addr:  addr,
	}
	router := &EngineGroup{
		engine:      engine,
		children:    make([]*EngineGroup, 0, 16),
		routerGroup: nil,
		path:        "",
		handlerHash: make(map[string][]localAPIHandler),
	}
	engine.gin.OPTIONS("/*all", func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.AbortWithStatus(http.StatusOK)
	})
	engine.gin.NoMethod(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.AbortWithStatus(http.StatusMethodNotAllowed)
	})
	engine.gin.NotFound404(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.AbortWithStatus(http.StatusNotFound)
	})
	apiLocalSupportRegister(addr, router)
	return router
}

func (this *EngineGroup) Group(path string) *EngineGroup {
	var routerGroup *gin.RouterGroup
	if this.routerGroup != nil {
		routerGroup = this.routerGroup.Group(path)
	} else {
		routerGroup = this.engine.gin.Group(path)
	}
	child := &EngineGroup{
		engine:      this.engine,
		routerGroup: routerGroup,
		path:        path,
		handlerHash: make(map[string][]localAPIHandler),
	}
	this.children = append(this.children, child)
	return child
}

const (
	METHOD_PUT    = iota
	METHOD_GET    = iota
	METHOD_POST   = iota
	METHOD_DELETE = iota
	method_COUNT  = iota
)

func (this *EngineGroup) register(method int, path string, handler APIHandler) {
	apiList, has := this.handlerHash[path]
	if !has {
		apiList = make([]localAPIHandler, method_COUNT)
		this.handlerHash[path] = apiList
	}
	apiList[method] = func(c *Context) {
		session, _ := this.engine.store.Get(c.Request, API_SESSION_NAME)
		handler(c, session)
		session.Save(c.Request, c.Writer)
	}

	handlerImp := this.getHandlerImp(handler)
	switch method {
	case METHOD_PUT:
		if this.routerGroup != nil {
			this.routerGroup.PUT(path, handlerImp)
		} else {
			this.engine.gin.PUT(path, handlerImp)
		}
	case METHOD_GET:
		if this.routerGroup != nil {
			this.routerGroup.GET(path, handlerImp)
		} else {
			this.engine.gin.GET(path, handlerImp)
		}
	case METHOD_POST:
		if this.routerGroup != nil {
			this.routerGroup.POST(path, handlerImp)
		} else {
			this.engine.gin.POST(path, handlerImp)
		}
	case METHOD_DELETE:
		if this.routerGroup != nil {
			this.routerGroup.DELETE(path, handlerImp)
		} else {
			this.engine.gin.DELETE(path, handlerImp)
		}
	}
}

func (this *EngineGroup) PUT(path string, handler APIHandler) {
	this.register(METHOD_PUT, path, handler)
}

func (this *EngineGroup) GET(path string, handler APIHandler) {
	this.register(METHOD_GET, path, handler)
}

func (this *EngineGroup) POST(path string, handler APIHandler) {
	this.register(METHOD_POST, path, handler)
}

func (this *EngineGroup) DELETE(path string, handler APIHandler) {
	this.register(METHOD_DELETE, path, handler)
}

func (this *EngineGroup) getHandlerImp(handler APIHandler) gin.HandlerFunc {
	return func(c *gin.Context) {
		var data interface{}
		body, _ := ioutil.ReadAll(c.Request.Body)
		err := json.Unmarshal(body, &data)
		if err != nil {
			var buffer bytes.Buffer
			buffer.Write(body)
			c.Request.Body = ioutil.NopCloser(bufio.NewReader(&buffer))
		}
		context := &Context{
			code:    unresolvedCode,
			ret:     nil,
			Data:    data,
			Params:  c.Params,
			QParams: c.Request.URL.Query(),
			Request: c.Request,
			Writer:  c.Writer,
		}
		session, _ := this.engine.store.Get(c.Request, API_SESSION_NAME)
		handler(context, session)
		session.Save(c.Request, c.Writer)
		if this.CORS {
			c.Writer.Header().Set("Access-Control-Allow-Origin", c.Request.Header.Get("Origin"))
	        c.Writer.Header().Set("Access-Control-Allow-Methods", "PUT, GET, POST, DELETE");
			c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		}
		switch context.code {
		case unresolvedCode:
			c.AbortWithStatus(http.StatusInternalServerError)
		case http.StatusMovedPermanently:
			c.Redirect(context.code, context.ret.(string))
		case http.StatusTemporaryRedirect:
			c.Redirect(context.code, context.ret.(string))
		default:
			switch context.ret.(type) {
			case nil:
				c.AbortWithStatus(context.code)
			case string:
				c.HTMLString(context.code, context.ret.(string))
				//c.String(context.code, context.ret.(string))
			default:
				c.JSON(context.code, context.ret)
			}
		}
	}
}

func (this *EngineGroup) Run() error {
	return this.engine.gin.Run(this.engine.addr)
}

func (this *EngineGroup) RunTLS(cert string, key string) error {
	return this.engine.gin.RunTLS(this.engine.addr, cert, key)
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"net"
	nethttp "net/http"
	"net/url"
	"strings"
	"time"
)

// Option is the function signature required to be considered an http.Option.
type Option func(*Transport) error

// WithTarget sets the outbound recipient of cloudevents when using an HTTP
// request.
func WithTarget(targetUrl string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http target option can not set nil transport")
		}
		targetUrl = strings.TrimSpace(targetUrl)
		if targetUrl != "" {
			var err error
			var target *url.URL
			target, err = url.Parse(targetUrl)
			if err != nil {
				return fmt.Errorf("http target option failed to parse target url: %s", err.Error())
			}

			if t.Req == nil {
				t.Req = &nethttp.Request{
					Method: nethttp.MethodPost,
				}
			}
			t.Req.URL = target
			return nil
		}
		return fmt.Errorf("http target option was empty string")
	}
}

// WithMethod sets the outbound recipient of cloudevents when using an HTTP
// request.
func WithMethod(method string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http method option can not set nil transport")
		}
		method = strings.TrimSpace(method)
		if method != "" {
			if t.Req == nil {
				t.Req = &nethttp.Request{}
			}
			t.Req.Method = method
			return nil
		}
		return fmt.Errorf("http method option was empty string")
	}
}

// WithHeader sets an additional default outbound header for all cloudevents
// when using an HTTP request.
func WithHeader(key, value string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http header option can not set nil transport")
		}
		key = strings.TrimSpace(key)
		if key != "" {
			if t.Req == nil {
				t.Req = &nethttp.Request{}
			}
			if t.Req.Header == nil {
				t.Req.Header = nethttp.Header{}
			}
			t.Req.Header.Add(key, value)
			return nil
		}
		return fmt.Errorf("http header option was empty string")
	}
}

// WithShutdownTimeout sets the shutdown timeout when the http server is being shutdown.
func WithShutdownTimeout(timeout time.Duration) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http shutdown timeout option can not set nil transport")
		}
		t.ShutdownTimeout = &timeout
		return nil
	}
}

// WithEncoding sets the encoding for clients with HTTP transports.
func WithEncoding(encoding Encoding) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http encoding option can not set nil transport")
		}
		t.Encoding = encoding
		return nil
	}
}

// WithDefaultEncodingSelector sets the encoding selection strategy for
// default encoding selections based on Event.
func WithDefaultEncodingSelector(fn EncodingSelector) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http default encoding selector option can not set nil transport")
		}
		if fn != nil {
			t.DefaultEncodingSelectionFn = fn
			return nil
		}
		return fmt.Errorf("http fn for DefaultEncodingSelector was nil")
	}
}

// WithContextBasedEncoding sets the encoding selection strategy for
// default encoding selections based context and then on Event, the encoded
// event will be the given version in the encoding specified by the given
// context, or Binary if not set.
func WithContextBasedEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http context based encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = ContextBasedEncodingSelectionStrategy
		return nil
	}
}

// WithBinaryEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Binary form.
func WithBinaryEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http binary encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultBinaryEncodingSelectionStrategy
		return nil
	}
}

// WithStructuredEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Structured form.
func WithStructuredEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http structured encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultStructuredEncodingSelectionStrategy
		return nil
	}
}

func checkListen(t *Transport, prefix string) error {
	switch {
	case t.Port != nil:
		return fmt.Errorf("%v port already set", prefix)
	case t.listener != nil:
		return fmt.Errorf("%v listener already set", prefix)
	}
	return nil
}

// WithPort sets the listening port for StartReceiver.
// Only one of WithListener  or WithPort is allowed.
func WithPort(port int) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http port option can not set nil transport")
		}
		if port < 0 || port > 65535 {
			return fmt.Errorf("http port option was given an invalid port: %d", port)
		}
		if err := checkListen(t, "http port option"); err != nil {
			return err
		}
		t.setPort(port)
		return nil
	}
}

// WithListener sets the listener for StartReceiver.
// Only one of WithListener or WithPort is allowed.
func WithListener(l net.Listener) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http listener option can not set nil transport")
		}
		if err := checkListen(t, "http port option"); err != nil {
			return err
		}
		t.listener = l
		_, err := t.listen()
		return err
	}
}

// WithPath sets the path to receive cloudevents on for HTTP transports.
func WithPath(path string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http path option can not set nil transport")
		}
		path = strings.TrimSpace(path)
		if len(path) == 0 {
			return fmt.Errorf("http path option was given an invalid path: %q", path)
		}
		t.Path = path
		return nil
	}
}

// Middleware is a function that takes an existing http.Handler and wraps it in middleware,
// returning the wrapped http.Handler.
type Middleware func(next nethttp.Handler) nethttp.Handler

// WithMiddleware adds an HTTP middleware to the transport. It may be specified multiple times.
// Middleware is applied to everything before it. For example
// `NewClient(WithMiddleware(foo), WithMiddleware(bar))` would result in `bar(foo(original))`.
func WithMiddleware(middleware Middleware) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http middleware option can not set nil transport")
		}
		t.middleware = append(t.middleware, middleware)
		return nil
	}
}

// WithLongPollTarget sets the receivers URL to perform long polling after
// StartReceiver is called.
func WithLongPollTarget(targetUrl string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http long poll target option can not set nil transport")
		}
		targetUrl = strings.TrimSpace(targetUrl)
		if targetUrl != "" {
			var err error
			var target *url.URL
			target, err = url.Parse(targetUrl)
			if err != nil {
				return fmt.Errorf("http long poll target option failed to parse target url: %s", err.Error())
			}

			if t.LongPollReq == nil {
				t.LongPollReq = &nethttp.Request{
					Method: nethttp.MethodGet,
				}
			}
			t.LongPollReq.URL = target
			return nil
		}
		return fmt.Errorf("http long poll target option was empty string")
	}
}
<commit_msg>Fix copy paste docs error (#307)<commit_after>package http

import (
	"fmt"
	"net"
	nethttp "net/http"
	"net/url"
	"strings"
	"time"
)

// Option is the function signature required to be considered an http.Option.
type Option func(*Transport) error

// WithTarget sets the outbound recipient of cloudevents when using an HTTP
// request.
func WithTarget(targetUrl string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http target option can not set nil transport")
		}
		targetUrl = strings.TrimSpace(targetUrl)
		if targetUrl != "" {
			var err error
			var target *url.URL
			target, err = url.Parse(targetUrl)
			if err != nil {
				return fmt.Errorf("http target option failed to parse target url: %s", err.Error())
			}

			if t.Req == nil {
				t.Req = &nethttp.Request{
					Method: nethttp.MethodPost,
				}
			}
			t.Req.URL = target
			return nil
		}
		return fmt.Errorf("http target option was empty string")
	}
}

// WithMethod sets the HTTP verb (GET, POST, PUT, etc.) to use
// when using an HTTP request.
func WithMethod(method string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http method option can not set nil transport")
		}
		method = strings.TrimSpace(method)
		if method != "" {
			if t.Req == nil {
				t.Req = &nethttp.Request{}
			}
			t.Req.Method = method
			return nil
		}
		return fmt.Errorf("http method option was empty string")
	}
}

// WithHeader sets an additional default outbound header for all cloudevents
// when using an HTTP request.
func WithHeader(key, value string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http header option can not set nil transport")
		}
		key = strings.TrimSpace(key)
		if key != "" {
			if t.Req == nil {
				t.Req = &nethttp.Request{}
			}
			if t.Req.Header == nil {
				t.Req.Header = nethttp.Header{}
			}
			t.Req.Header.Add(key, value)
			return nil
		}
		return fmt.Errorf("http header option was empty string")
	}
}

// WithShutdownTimeout sets the shutdown timeout when the http server is being shutdown.
func WithShutdownTimeout(timeout time.Duration) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http shutdown timeout option can not set nil transport")
		}
		t.ShutdownTimeout = &timeout
		return nil
	}
}

// WithEncoding sets the encoding for clients with HTTP transports.
func WithEncoding(encoding Encoding) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http encoding option can not set nil transport")
		}
		t.Encoding = encoding
		return nil
	}
}

// WithDefaultEncodingSelector sets the encoding selection strategy for
// default encoding selections based on Event.
func WithDefaultEncodingSelector(fn EncodingSelector) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http default encoding selector option can not set nil transport")
		}
		if fn != nil {
			t.DefaultEncodingSelectionFn = fn
			return nil
		}
		return fmt.Errorf("http fn for DefaultEncodingSelector was nil")
	}
}

// WithContextBasedEncoding sets the encoding selection strategy for
// default encoding selections based context and then on Event, the encoded
// event will be the given version in the encoding specified by the given
// context, or Binary if not set.
func WithContextBasedEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http context based encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = ContextBasedEncodingSelectionStrategy
		return nil
	}
}

// WithBinaryEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Binary form.
func WithBinaryEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http binary encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultBinaryEncodingSelectionStrategy
		return nil
	}
}

// WithStructuredEncoding sets the encoding selection strategy for
// default encoding selections based on Event, the encoded event will be the
// given version in Structured form.
func WithStructuredEncoding() Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http structured encoding option can not set nil transport")
		}

		t.DefaultEncodingSelectionFn = DefaultStructuredEncodingSelectionStrategy
		return nil
	}
}

func checkListen(t *Transport, prefix string) error {
	switch {
	case t.Port != nil:
		return fmt.Errorf("%v port already set", prefix)
	case t.listener != nil:
		return fmt.Errorf("%v listener already set", prefix)
	}
	return nil
}

// WithPort sets the listening port for StartReceiver.
// Only one of WithListener  or WithPort is allowed.
func WithPort(port int) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http port option can not set nil transport")
		}
		if port < 0 || port > 65535 {
			return fmt.Errorf("http port option was given an invalid port: %d", port)
		}
		if err := checkListen(t, "http port option"); err != nil {
			return err
		}
		t.setPort(port)
		return nil
	}
}

// WithListener sets the listener for StartReceiver.
// Only one of WithListener or WithPort is allowed.
func WithListener(l net.Listener) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http listener option can not set nil transport")
		}
		if err := checkListen(t, "http port option"); err != nil {
			return err
		}
		t.listener = l
		_, err := t.listen()
		return err
	}
}

// WithPath sets the path to receive cloudevents on for HTTP transports.
func WithPath(path string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http path option can not set nil transport")
		}
		path = strings.TrimSpace(path)
		if len(path) == 0 {
			return fmt.Errorf("http path option was given an invalid path: %q", path)
		}
		t.Path = path
		return nil
	}
}

// Middleware is a function that takes an existing http.Handler and wraps it in middleware,
// returning the wrapped http.Handler.
type Middleware func(next nethttp.Handler) nethttp.Handler

// WithMiddleware adds an HTTP middleware to the transport. It may be specified multiple times.
// Middleware is applied to everything before it. For example
// `NewClient(WithMiddleware(foo), WithMiddleware(bar))` would result in `bar(foo(original))`.
func WithMiddleware(middleware Middleware) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http middleware option can not set nil transport")
		}
		t.middleware = append(t.middleware, middleware)
		return nil
	}
}

// WithLongPollTarget sets the receivers URL to perform long polling after
// StartReceiver is called.
func WithLongPollTarget(targetUrl string) Option {
	return func(t *Transport) error {
		if t == nil {
			return fmt.Errorf("http long poll target option can not set nil transport")
		}
		targetUrl = strings.TrimSpace(targetUrl)
		if targetUrl != "" {
			var err error
			var target *url.URL
			target, err = url.Parse(targetUrl)
			if err != nil {
				return fmt.Errorf("http long poll target option failed to parse target url: %s", err.Error())
			}

			if t.LongPollReq == nil {
				t.LongPollReq = &nethttp.Request{
					Method: nethttp.MethodGet,
				}
			}
			t.LongPollReq.URL = target
			return nil
		}
		return fmt.Errorf("http long poll target option was empty string")
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 Google LLC

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 review

import (
	"encoding/base64"
	"fmt"

	"github.com/golang/glog"
	"github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/attestation"
	"github.com/grafeas/kritis/pkg/kritis/container"
	"github.com/grafeas/kritis/pkg/kritis/metadata"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
)

// ValidatingTransport allows the caller to obtain validated attestations for a given container image.
// Implementations should return trusted and verified attestations.
type ValidatingTransport interface {
	GetValidatedAttestations(image string) ([]attestation.ValidatedAttestation, error)
}

// Implements ValidatingTransport.
type AttestorValidatingTransport struct {
	Client   metadata.ReadOnlyClient
	Attestor v1beta1.AttestationAuthority
}

func (avt *AttestorValidatingTransport) GetValidatedAttestations(image string) ([]attestation.ValidatedAttestation, error) {
	keys := map[string]string{}
	for _, keyData := range avt.Attestor.Spec.PublicKeyList {
		key, fingerprint, err := secrets.KeyAndFingerprint(keyData)
		if err != nil {
			glog.Warningf("Error parsing key for %q: %v", avt.Attestor.Name, err)
		} else {
			if _, ok := keys[fingerprint]; ok {
				glog.Warningf("Duplicate keys with fingerprint %s for %q.", fingerprint, avt.Attestor.Name)
			}
			keys[fingerprint] = key
		}
	}
	if len(keys) == 0 {
		return nil, fmt.Errorf("Unable to find any valid key for %q", avt.Attestor.Name)
	}

	out := []attestation.ValidatedAttestation{}
	host, err := container.NewAtomicContainerSig(image, map[string]string{})
	if err != nil {
		glog.Error(err)
		return nil, err
	}
	attestations, err := avt.Client.Attestations(image, &avt.Attestor)
	if err != nil {
		glog.Error(err)
		return nil, err
	}
	for _, a := range attestations {
		// TODO(acamadeo): Temporarily hardcoding this to work only for PGP Attestations.
		signature := a.PgpAttestation.Signature.Signature
		decodedSig := make([]byte, base64.StdEncoding.DecodedLen(len(signature)))
		_, err := base64.StdEncoding.Decode(decodedSig, signature)
		if err != nil {
			glog.Infof("Cannot base64 decode signature: %v", err)
			continue
		}
		keyId := a.PgpAttestation.Signature.PublicKeyId
		if err = host.VerifyAttestationSignature(keys[keyId], string(decodedSig)); err != nil {
			glog.Infof("Could not find or verify attestation for attestor %s: %s", keyId, err.Error())
			continue
		}
		out = append(out, attestation.ValidatedAttestation{AttestorName: avt.Attestor.Name, Image: image})
	}
	return out, nil
}
<commit_msg>Revert base64.StdEncoding.Decode<commit_after>/*
Copyright 2019 Google LLC

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 review

import (
	"encoding/base64"
	"fmt"

	"github.com/golang/glog"
	"github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/attestation"
	"github.com/grafeas/kritis/pkg/kritis/container"
	"github.com/grafeas/kritis/pkg/kritis/metadata"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
)

// ValidatingTransport allows the caller to obtain validated attestations for a given container image.
// Implementations should return trusted and verified attestations.
type ValidatingTransport interface {
	GetValidatedAttestations(image string) ([]attestation.ValidatedAttestation, error)
}

// Implements ValidatingTransport.
type AttestorValidatingTransport struct {
	Client   metadata.ReadOnlyClient
	Attestor v1beta1.AttestationAuthority
}

func (avt *AttestorValidatingTransport) GetValidatedAttestations(image string) ([]attestation.ValidatedAttestation, error) {
	keys := map[string]string{}
	for _, keyData := range avt.Attestor.Spec.PublicKeyList {
		key, fingerprint, err := secrets.KeyAndFingerprint(keyData)
		if err != nil {
			glog.Warningf("Error parsing key for %q: %v", avt.Attestor.Name, err)
		} else {
			if _, ok := keys[fingerprint]; ok {
				glog.Warningf("Duplicate keys with fingerprint %s for %q.", fingerprint, avt.Attestor.Name)
			}
			keys[fingerprint] = key
		}
	}
	if len(keys) == 0 {
		return nil, fmt.Errorf("Unable to find any valid key for %q", avt.Attestor.Name)
	}

	out := []attestation.ValidatedAttestation{}
	host, err := container.NewAtomicContainerSig(image, map[string]string{})
	if err != nil {
		glog.Error(err)
		return nil, err
	}
	attestations, err := avt.Client.Attestations(image, &avt.Attestor)
	if err != nil {
		glog.Error(err)
		return nil, err
	}
	for _, a := range attestations {
		// TODO(acamadeo): Temporarily hardcoding this to work only for PGP Attestations.
		signature := a.PgpAttestation.Signature.Signature
		decodedSig, err := base64.StdEncoding.DecodeString(string(signature))
		if err != nil {
			glog.Infof("Cannot base64 decode signature: %v", err)
			continue
		}
		keyId := a.PgpAttestation.Signature.PublicKeyId
		if err = host.VerifyAttestationSignature(keys[keyId], string(decodedSig)); err != nil {
			glog.Infof("Could not find or verify attestation for attestor %s: %s", keyId, err.Error())
			continue
		}
		out = append(out, attestation.ValidatedAttestation{AttestorName: avt.Attestor.Name, Image: image})
	}
	return out, nil
}
<|endoftext|>