<commit_before>package tchannel

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

import (
	"fmt"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

func testHandlerFunc(t *testing.T, f func(t *testing.T, ctx context.Context, call *InboundCall)) Handler {
	return HandlerFunc(func(ctx context.Context, call *InboundCall) {
		f(t, ctx, call)
	})
}

func serverBusy(t *testing.T, ctx context.Context, call *InboundCall) {
	call.Response().SendSystemError(ErrServerBusy)
}

func timeout(t *testing.T, ctx context.Context, call *InboundCall) {
	deadline, _ := ctx.Deadline()
	time.Sleep(deadline.Add(time.Second * 1).Sub(time.Now()))
	echo(t, ctx, call)
}

func echo(t *testing.T, ctx context.Context, call *InboundCall) {
	fmt.Println("Entering echo")
	var inArg2 BytesInput
	var inArg3 BytesInput

	require.NoError(t, call.ReadArg2(&inArg2))
	require.NoError(t, call.ReadArg3(&inArg3))
	require.NoError(t, call.Response().WriteArg2(BytesOutput(inArg2)))
	require.NoError(t, call.Response().WriteArg3(BytesOutput(inArg3)))
}

func TestRoundTrip(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {

		ch.Register(testHandlerFunc(t, echo), "Capture", "ping")

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		call, err := ch.BeginCall(ctx, hostPort, "Capture", "ping")
		require.NoError(t, err)

		require.NoError(t, call.WriteArg2(BytesOutput("Hello Header")))
		require.NoError(t, call.WriteArg3(BytesOutput("Body Sent")))

		var respArg2 BytesInput
		require.NoError(t, call.Response().ReadArg2(&respArg2))
		assert.Equal(t, []byte("Hello Header"), []byte(respArg2))

		var respArg3 BytesInput
		require.NoError(t, call.Response().ReadArg3(&respArg3))
		assert.Equal(t, []byte("Body Sent"), []byte(respArg3))
	})
}

func TestBadRequest(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		_, _, err := sendRecv(ctx, ch, hostPort, "Nowhere", "Noone", []byte("Headers"), []byte("Body"))
		require.NotNil(t, err)
		assert.Equal(t, ErrCodeBadRequest, GetSystemErrorCode(err))
	})
}

func TestServerBusy(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ch.Register(testHandlerFunc(t, serverBusy), "TestService", "busy")

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		_, _, err := sendRecv(ctx, ch, hostPort, "TestService", "busy", []byte("Arg2"), []byte("Arg3"))
		require.NotNil(t, err)
		assert.Equal(t, ErrCodeBusy, GetSystemErrorCode(err))
	})
}

func TestTimeout(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ch.Register(testHandlerFunc(t, timeout), "TestService", "timeout")

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
		defer cancel()

		_, _, err := sendRecv(ctx, ch, hostPort, "TestService", "timeout", []byte("Arg2"), []byte("Arg3"))

		// TODO(mmihic): Maybe translate this into ErrTimeout (or vice versa)?
		assert.Equal(t, context.DeadlineExceeded, err)
	})
}

func testFragmentation(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ch.Register(testHandlerFunc(t, echo), "TestService", "echo")

		arg2 := make([]byte, MaxFramePayloadSize*2)
		for i := 0; i < len(arg2); i++ {
			arg2[i] = byte('a' + (i % 10))
		}

		arg3 := make([]byte, MaxFramePayloadSize*3)
		for i := 0; i < len(arg3); i++ {
			arg3[i] = byte('A' + (i % 10))
		}

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
		defer cancel()

		respArg2, respArg3, err := sendRecv(ctx, ch, hostPort, "TestService", "echo", arg2, arg3)
		require.NoError(t, err)
		assert.Equal(t, arg2, respArg2)
		assert.Equal(t, arg3, respArg3)
	})
}

func sendRecv(ctx context.Context, ch *Channel, hostPort string, serviceName, operation string,
	arg2, arg3 []byte) ([]byte, []byte, error) {

	call, err := ch.BeginCall(ctx, hostPort, serviceName, operation)
	if err != nil {
		return nil, nil, err
	}

	if err := call.WriteArg2(BytesOutput(arg2)); err != nil {
		return nil, nil, err
	}

	if err := call.WriteArg3(BytesOutput(arg3)); err != nil {
		return nil, nil, err
	}

	var respArg2 BytesInput
	if err := call.Response().ReadArg2(&respArg2); err != nil {
		return nil, nil, err
	}

	var respArg3 BytesInput
	if err := call.Response().ReadArg3(&respArg3); err != nil {
		return nil, nil, err
	}

	return []byte(respArg2), []byte(respArg3), nil
}

func withTestChannel(t *testing.T, f func(ch *Channel, hostPort string)) {
	opts := ChannelOptions{
		Logger: SimpleLogger,
	}

	ch, err := NewChannel(&opts)
	require.Nil(t, err)

	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.Nil(t, err)

	go ch.Serve(l)

	f(ch, l.Addr().String())
}
<commit_msg>Re-enable fragmentation test<commit_after>package tchannel

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

import (
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

func testHandlerFunc(t *testing.T, f func(t *testing.T, ctx context.Context, call *InboundCall)) Handler {
	return HandlerFunc(func(ctx context.Context, call *InboundCall) {
		f(t, ctx, call)
	})
}

func serverBusy(t *testing.T, ctx context.Context, call *InboundCall) {
	call.Response().SendSystemError(ErrServerBusy)
}

func timeout(t *testing.T, ctx context.Context, call *InboundCall) {
	deadline, _ := ctx.Deadline()
	time.Sleep(deadline.Add(time.Second * 1).Sub(time.Now()))
	echo(t, ctx, call)
}

func echo(t *testing.T, ctx context.Context, call *InboundCall) {
	var inArg2 BytesInput
	var inArg3 BytesInput

	require.NoError(t, call.ReadArg2(&inArg2))
	require.NoError(t, call.ReadArg3(&inArg3))
	require.NoError(t, call.Response().WriteArg2(BytesOutput(inArg2)))
	require.NoError(t, call.Response().WriteArg3(BytesOutput(inArg3)))
}

func TestRoundTrip(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {

		ch.Register(testHandlerFunc(t, echo), "Capture", "ping")

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		call, err := ch.BeginCall(ctx, hostPort, "Capture", "ping")
		require.NoError(t, err)

		require.NoError(t, call.WriteArg2(BytesOutput("Hello Header")))
		require.NoError(t, call.WriteArg3(BytesOutput("Body Sent")))

		var respArg2 BytesInput
		require.NoError(t, call.Response().ReadArg2(&respArg2))
		assert.Equal(t, []byte("Hello Header"), []byte(respArg2))

		var respArg3 BytesInput
		require.NoError(t, call.Response().ReadArg3(&respArg3))
		assert.Equal(t, []byte("Body Sent"), []byte(respArg3))
	})
}

func TestBadRequest(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		_, _, err := sendRecv(ctx, ch, hostPort, "Nowhere", "Noone", []byte("Headers"), []byte("Body"))
		require.NotNil(t, err)
		assert.Equal(t, ErrCodeBadRequest, GetSystemErrorCode(err))
	})
}

func TestServerBusy(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ch.Register(testHandlerFunc(t, serverBusy), "TestService", "busy")

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
		defer cancel()

		_, _, err := sendRecv(ctx, ch, hostPort, "TestService", "busy", []byte("Arg2"), []byte("Arg3"))
		require.NotNil(t, err)
		assert.Equal(t, ErrCodeBusy, GetSystemErrorCode(err))
	})
}

func TestTimeout(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ch.Register(testHandlerFunc(t, timeout), "TestService", "timeout")

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
		defer cancel()

		_, _, err := sendRecv(ctx, ch, hostPort, "TestService", "timeout", []byte("Arg2"), []byte("Arg3"))

		// TODO(mmihic): Maybe translate this into ErrTimeout (or vice versa)?
		assert.Equal(t, context.DeadlineExceeded, err)
	})
}

func TestFragmentation(t *testing.T) {
	withTestChannel(t, func(ch *Channel, hostPort string) {
		ch.Register(testHandlerFunc(t, echo), "TestService", "echo")

		arg2 := make([]byte, MaxFramePayloadSize*2)
		for i := 0; i < len(arg2); i++ {
			arg2[i] = byte('a' + (i % 10))
		}

		arg3 := make([]byte, MaxFramePayloadSize*3)
		for i := 0; i < len(arg3); i++ {
			arg3[i] = byte('A' + (i % 10))
		}

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
		defer cancel()

		respArg2, respArg3, err := sendRecv(ctx, ch, hostPort, "TestService", "echo", arg2, arg3)
		require.NoError(t, err)
		assert.Equal(t, arg2, respArg2)
		assert.Equal(t, arg3, respArg3)
	})
}

func sendRecv(ctx context.Context, ch *Channel, hostPort string, serviceName, operation string,
	arg2, arg3 []byte) ([]byte, []byte, error) {

	call, err := ch.BeginCall(ctx, hostPort, serviceName, operation)
	if err != nil {
		return nil, nil, err
	}

	if err := call.WriteArg2(BytesOutput(arg2)); err != nil {
		return nil, nil, err
	}

	if err := call.WriteArg3(BytesOutput(arg3)); err != nil {
		return nil, nil, err
	}

	var respArg2 BytesInput
	if err := call.Response().ReadArg2(&respArg2); err != nil {
		return nil, nil, err
	}

	var respArg3 BytesInput
	if err := call.Response().ReadArg3(&respArg3); err != nil {
		return nil, nil, err
	}

	return []byte(respArg2), []byte(respArg3), nil
}

func withTestChannel(t *testing.T, f func(ch *Channel, hostPort string)) {
	opts := ChannelOptions{
		Logger: SimpleLogger,
	}

	ch, err := NewChannel(&opts)
	require.Nil(t, err)

	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.Nil(t, err)

	go ch.Serve(l)

	f(ch, l.Addr().String())
}
<|endoftext|>


<commit_before>package sgload

import (
	"encoding/json"
	"fmt"
	"sync"
)

type ReadLoadRunner struct {
	LoadRunner
	ReadLoadSpec ReadLoadSpec
}

func NewReadLoadRunner(rls ReadLoadSpec) *ReadLoadRunner {

	rls.MustValidate()

	loadRunner := LoadRunner{
		LoadSpec: rls.LoadSpec,
	}
	loadRunner.CreateStatsdClient()

	return &ReadLoadRunner{
		LoadRunner:   loadRunner,
		ReadLoadSpec: rls,
	}

}

func (rlr ReadLoadRunner) Run() error {

	// Create a wait group to see when all the reader goroutines have finished
	var wg sync.WaitGroup

	// Create writer goroutines
	readers, err := rlr.createReaders(&wg)
	if err != nil {
		return fmt.Errorf("Error creating readers: %v", err)
	}
	for _, reader := range readers {
		go reader.Run()
	}

	// block until readers are done
	logger.Info("Waiting for readers to finish")
	wg.Wait()
	logger.Info("Readers finished")

	return nil

}

func (rlr ReadLoadRunner) createReaders(wg *sync.WaitGroup) ([]*Reader, error) {

	readers := []*Reader{}
	var userCreds []UserCred
	var err error

	switch rlr.ReadLoadSpec.CreateReaders {
	case true:
		userCreds = rlr.generateUserCreds()
	default:
		userCreds, err = rlr.loadUserCredsFromArgs()
		if err != nil {
			return readers, fmt.Errorf("Error loading user creds from args: %v", err)
		}
	}

	for userId := 0; userId < rlr.ReadLoadSpec.NumReaders; userId++ {
		userCred := userCreds[userId]
		dataStore := rlr.createDataStore()
		dataStore.SetUserCreds(userCred)

		// get channels that should be assigned to this reader
		sgChannels := rlr.assignChannelsToReader(rlr.generateChannelNames())

		reader := NewReader(
			wg,
			userId,
			userCred,
			dataStore,
			rlr.ReadLoadSpec.BatchSize,
		)
		reader.SetChannels(sgChannels)
		reader.SetBatchSize(rlr.ReadLoadSpec.BatchSize)
		reader.SetNumDocsExpected(rlr.numDocsExpectedPerReader())
		reader.CreateDataStoreUser = rlr.ReadLoadSpec.CreateReaders
		readers = append(readers, reader)
		wg.Add(1)
	}

	return readers, nil
}

// TODO: duplicated code with WriteLoadRunner.loadUserCredsFromArgs()
func (rlr ReadLoadRunner) loadUserCredsFromArgs() ([]UserCred, error) {

	userCreds := []UserCred{}
	var err error

	switch {
	case rlr.ReadLoadSpec.ReaderCreds != "":
		logger.Info("Load writer creds from CLI args")
		err = json.Unmarshal([]byte(rlr.ReadLoadSpec.ReaderCreds), &userCreds)
		if err != nil {
			return userCreds, err
		}
		for _, userCred := range userCreds {
			if userCred.Empty() {
				return userCreds, fmt.Errorf("User credentials empty: %+v", userCred)
			}
		}
	case rlr.ReadLoadSpec.TestSessionID != "" && rlr.ReadLoadSpec.DidAutoGenTestSessionID == false:
		// If the user explicitly provided a test session ID, then use that
		// to generate user credentials to use.  Presumably these credentials
		// were created before in previous runs.  Doesn't make sense to use
		// this with auto-generated test session ID's, since there is no way
		// that the Sync Gateway will have those users created from prev. runs
		logger.Info("Generate user creds from test session id")
		userCreds = rlr.generateUserCreds()
	default:
		return userCreds, fmt.Errorf("You need to either create writers, specify a test session ID, or specify writer user credentials.  See CLI help.")

	}

	if len(userCreds) != rlr.ReadLoadSpec.NumReaders {
		return userCreds, fmt.Errorf("You only provided %d user credentials, but specified %d readers", len(userCreds), rlr.ReadLoadSpec.NumReaders)
	}

	return userCreds, err
}

// Calculate how many docs each reader is expected to pull.  Find out how many docs are
// in each channel, and then find out how many channels each reader is pulling from,
// and then multiply to get the number docs each reader is expected to pull.
func (rlr ReadLoadRunner) numDocsExpectedPerReader() int {

	numDocsPerChannel := rlr.ReadLoadSpec.NumDocs / rlr.ReadLoadSpec.NumChannels
	docsPerReader := numDocsPerChannel * rlr.ReadLoadSpec.NumChansPerReader

	return docsPerReader

}

// Given the full list of SG channel names for this scenario, assign one more more
// SG channels to this particular reader.  This means that when the reader user is
// created, this will have these channels listed in their admin_channels field
// so they pull these channels when hittting the _changes feed.
func (rlr ReadLoadRunner) assignChannelsToReader(sgChannels []string) []string {

	assignedChannels := []string{}

	if rlr.ReadLoadSpec.NumChansPerReader > len(sgChannels) {
		panic(fmt.Sprintf("Cannot have more chans per reader (%d) than total channels (%d)", rlr.ReadLoadSpec.NumChansPerReader, len(sgChannels)))
	}

	for i := 0; i < rlr.ReadLoadSpec.NumChansPerReader; i++ {
		sgChannel := sgChannels[i]
		assignedChannels = append(assignedChannels, sgChannel)
	}

	return assignedChannels

}

func (rlr ReadLoadRunner) generateUserCreds() []UserCred {
	return rlr.LoadRunner.generateUserCreds(rlr.ReadLoadSpec.NumReaders, "readload")
}
<commit_msg>add logging around rlr.numDocsExpectedPerReader()<commit_after>package sgload

import (
	"encoding/json"
	"fmt"
	"sync"
)

type ReadLoadRunner struct {
	LoadRunner
	ReadLoadSpec ReadLoadSpec
}

func NewReadLoadRunner(rls ReadLoadSpec) *ReadLoadRunner {

	rls.MustValidate()

	loadRunner := LoadRunner{
		LoadSpec: rls.LoadSpec,
	}
	loadRunner.CreateStatsdClient()

	return &ReadLoadRunner{
		LoadRunner:   loadRunner,
		ReadLoadSpec: rls,
	}

}

func (rlr ReadLoadRunner) Run() error {

	// Create a wait group to see when all the reader goroutines have finished
	var wg sync.WaitGroup

	// Create writer goroutines
	readers, err := rlr.createReaders(&wg)
	if err != nil {
		return fmt.Errorf("Error creating readers: %v", err)
	}
	for _, reader := range readers {
		go reader.Run()
	}

	// block until readers are done
	logger.Info("Waiting for readers to finish")
	wg.Wait()
	logger.Info("Readers finished")

	return nil

}

func (rlr ReadLoadRunner) createReaders(wg *sync.WaitGroup) ([]*Reader, error) {

	readers := []*Reader{}
	var userCreds []UserCred
	var err error
	numDocsExpectedPerReader := rlr.numDocsExpectedPerReader()

	switch rlr.ReadLoadSpec.CreateReaders {
	case true:
		userCreds = rlr.generateUserCreds()
	default:
		userCreds, err = rlr.loadUserCredsFromArgs()
		if err != nil {
			return readers, fmt.Errorf("Error loading user creds from args: %v", err)
		}
	}

	for userId := 0; userId < rlr.ReadLoadSpec.NumReaders; userId++ {
		userCred := userCreds[userId]
		dataStore := rlr.createDataStore()
		dataStore.SetUserCreds(userCred)

		// get channels that should be assigned to this reader
		sgChannels := rlr.assignChannelsToReader(rlr.generateChannelNames())

		reader := NewReader(
			wg,
			userId,
			userCred,
			dataStore,
			rlr.ReadLoadSpec.BatchSize,
		)
		reader.SetChannels(sgChannels)
		reader.SetBatchSize(rlr.ReadLoadSpec.BatchSize)
		reader.SetNumDocsExpected(numDocsExpectedPerReader)
		reader.CreateDataStoreUser = rlr.ReadLoadSpec.CreateReaders
		readers = append(readers, reader)
		wg.Add(1)
	}

	return readers, nil
}

// TODO: duplicated code with WriteLoadRunner.loadUserCredsFromArgs()
func (rlr ReadLoadRunner) loadUserCredsFromArgs() ([]UserCred, error) {

	userCreds := []UserCred{}
	var err error

	switch {
	case rlr.ReadLoadSpec.ReaderCreds != "":
		logger.Info("Load writer creds from CLI args")
		err = json.Unmarshal([]byte(rlr.ReadLoadSpec.ReaderCreds), &userCreds)
		if err != nil {
			return userCreds, err
		}
		for _, userCred := range userCreds {
			if userCred.Empty() {
				return userCreds, fmt.Errorf("User credentials empty: %+v", userCred)
			}
		}
	case rlr.ReadLoadSpec.TestSessionID != "" && rlr.ReadLoadSpec.DidAutoGenTestSessionID == false:
		// If the user explicitly provided a test session ID, then use that
		// to generate user credentials to use.  Presumably these credentials
		// were created before in previous runs.  Doesn't make sense to use
		// this with auto-generated test session ID's, since there is no way
		// that the Sync Gateway will have those users created from prev. runs
		logger.Info("Generate user creds from test session id")
		userCreds = rlr.generateUserCreds()
	default:
		return userCreds, fmt.Errorf("You need to either create writers, specify a test session ID, or specify writer user credentials.  See CLI help.")

	}

	if len(userCreds) != rlr.ReadLoadSpec.NumReaders {
		return userCreds, fmt.Errorf("You only provided %d user credentials, but specified %d readers", len(userCreds), rlr.ReadLoadSpec.NumReaders)
	}

	return userCreds, err
}

// Calculate how many docs each reader is expected to pull.  Find out how many docs are
// in each channel, and then find out how many channels each reader is pulling from,
// and then multiply to get the number docs each reader is expected to pull.
func (rlr ReadLoadRunner) numDocsExpectedPerReader() int {

	numDocsPerChannel := rlr.ReadLoadSpec.NumDocs / rlr.ReadLoadSpec.NumChannels

	logger.Info("Calculated numDocsPerChannel", "numDocsPerChannel", numDocsPerChannel, "numDocs", rlr.ReadLoadSpec.NumDocs, "numChannels", rlr.ReadLoadSpec.NumChannels)

	docsPerReader := numDocsPerChannel * rlr.ReadLoadSpec.NumChansPerReader

	logger.Info("Calculated docsPerReader", "docsPerReader", docsPerReader, "numDocsPerChannel", numDocsPerChannel, "numChansPerReader", rlr.ReadLoadSpec.NumChansPerReader)

	return docsPerReader

}

// Given the full list of SG channel names for this scenario, assign one more more
// SG channels to this particular reader.  This means that when the reader user is
// created, this will have these channels listed in their admin_channels field
// so they pull these channels when hittting the _changes feed.
func (rlr ReadLoadRunner) assignChannelsToReader(sgChannels []string) []string {

	assignedChannels := []string{}

	if rlr.ReadLoadSpec.NumChansPerReader > len(sgChannels) {
		panic(fmt.Sprintf("Cannot have more chans per reader (%d) than total channels (%d)", rlr.ReadLoadSpec.NumChansPerReader, len(sgChannels)))
	}

	for i := 0; i < rlr.ReadLoadSpec.NumChansPerReader; i++ {
		sgChannel := sgChannels[i]
		assignedChannels = append(assignedChannels, sgChannel)
	}

	return assignedChannels

}

func (rlr ReadLoadRunner) generateUserCreds() []UserCred {
	return rlr.LoadRunner.generateUserCreds(rlr.ReadLoadSpec.NumReaders, "readload")
}
<|endoftext|>


<commit_before>package http

/*
现实思路:
    分为三种类型url
	1.普通的:如/fff/ddd/lll
	2.有映射值的:如/user/{who}/info
	3.尾部全部匹配的:如/user/*('*'只可以用于尾部)
    先进行分组,将普通的于要进行取值的分开
    普通的对于一个map,直接使用map[string]获取handlerFunc
    要取值得将url利用"/"切分成数组,匹配是再将实际url切分成数组进行比对并取值

*/

import (
	"fmt"
	"log"
	"net/http"
	"strings"

	"github.com/fudali113/golib/utils"
)

const (
	ALL                  = "*"
	URL_PARA_PREFIX_FLAG = "{"
	URL_PARA_LAST_FLAG   = "}"
	URL_PARA_FLAG        = "{}"
	EMPTY                = ""
)

//handler与filter容器
var (
	handlerMap *handleFuncMap
	filters    []Filter
)

func Start(port int) {
	log.Printf("server is starting , listen port is %d", port)
	err := http.ListenAndServe(fmt.Sprintf(":%d", port), &DoobHandler{filters: filters, handlerMap: handlerMap})
	if err != nil {
		log.Printf("start is fail => %s", err.Error())
	}
}

/**
 * 注册一个handler
 */
func AddHandlerFunc(url string, methodStr string, handler http.HandlerFunc) {
	paras := utils.Split(url, "/")
	matchParaCount := 0
	urlinfo := &urlInfo{}
	restHandler := &restHandlerFunc{function: handler, methodStr: strings.ToLower(methodStr)}
	for i, v := range paras {
		matchParaCount++
		para := strings.TrimSpace(v)
		//para = strings.ToLower(para)
		if para[0] == URL_PARA_PREFIX_FLAG[0] && para[len(para)-1] == URL_PARA_LAST_FLAG[0] {
			urlinfo.addUrlPara(urlMacthPara{urlPara: para[1 : len(para)-1], matchInfo: URL_PARA_FLAG})
		} else if para == ALL {
			if i == len(paras)-1 {
				urlinfo.addUrlPara(urlMacthPara{urlPara: para, matchInfo: ALL})
				handlerMap.lastAllMatch[url[:len(url)-1]] = restHandler
			} else {
				urlinfo.addUrlPara(urlMacthPara{urlPara: para, matchInfo: URL_PARA_FLAG})
			}
		} else {
			matchParaCount--
			urlinfo.addUrlPara(urlMacthPara{urlPara: para, matchInfo: EMPTY})
		}
	}
	if matchParaCount == 0 {
		handlerMap.simple[url] = restHandler
	} else {
		urlinfo.handler = restHandler
		len := urlinfo.len()
		handlerMap.rest.urls[len] = append(handlerMap.rest.urls[len], urlinfo)
	}
}

/**
 * 添加一个过滤器
 */
func AddFilter(f Filter) {
	filters = append(filters, f)
}

/**
 * 实现http接口
 */
func (this *DoobHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	for i := range this.filters {
		if this.filters[i].Filter(res, req) {
			continue
		} else {
			return
		}
	}
	handler, err := this.handlerMap.getHandler(req)
	if err != nil {
		fmt.Println(err)
		errStr := err.Error()
		if strings.Index(errStr, "method not match") >= 0 {
			res.WriteHeader(405)
		} else {
			res.WriteHeader(404)
		}
		return
	}
	handler(res, req)
}

func init() {
	simple := map[string]*restHandlerFunc{}
	rest := &restHandlerMap{urls: map[int][]*urlInfo{}}
	last := map[string]*restHandlerFunc{}
	handlerMap = &handleFuncMap{
		simple:       simple,
		rest:         rest,
		lastAllMatch: last,
	}
	filters = []Filter{}
	AddHandlerFunc("/", "", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("doob "))
	})
}
<commit_msg>添加一次处理计时，添加没匹配的路由的日志<commit_after>package http

/*
现实思路:
    分为三种类型url
	1.普通的:如/fff/ddd/lll
	2.有映射值的:如/user/{who}/info
	3.尾部全部匹配的:如/user/*('*'只可以用于尾部)
    先进行分组,将普通的于要进行取值的分开
    普通的对于一个map,直接使用map[string]获取handlerFunc
    要取值得将url利用"/"切分成数组,匹配是再将实际url切分成数组进行比对并取值

*/

import (
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/fudali113/golib/utils"
)

const (
	ALL                  = "*"
	URL_PARA_PREFIX_FLAG = "{"
	URL_PARA_LAST_FLAG   = "}"
	URL_PARA_FLAG        = "{}"
	EMPTY                = ""
)

//handler与filter容器
var (
	handlerMap *handleFuncMap
	filters    []Filter
)

func Start(port int) {
	log.Printf("server is starting , listen port is %d", port)
	err := http.ListenAndServe(fmt.Sprintf(":%d", port), &DoobHandler{filters: filters, handlerMap: handlerMap})
	if err != nil {
		log.Printf("start is fail => %s", err.Error())
	}
}

/**
 * 注册一个handler
 */
func AddHandlerFunc(url string, methodStr string, handler http.HandlerFunc) {
	paras := utils.Split(url, "/")
	matchParaCount := 0
	urlinfo := &urlInfo{}
	restHandler := &restHandlerFunc{function: handler, methodStr: strings.ToLower(methodStr)}
	for i, v := range paras {
		matchParaCount++
		para := strings.TrimSpace(v)
		//para = strings.ToLower(para)
		if para[0] == URL_PARA_PREFIX_FLAG[0] && para[len(para)-1] == URL_PARA_LAST_FLAG[0] {
			urlinfo.addUrlPara(urlMacthPara{urlPara: para[1 : len(para)-1], matchInfo: URL_PARA_FLAG})
		} else if para == ALL {
			if i == len(paras)-1 {
				urlinfo.addUrlPara(urlMacthPara{urlPara: para, matchInfo: ALL})
				handlerMap.lastAllMatch[url[:len(url)-1]] = restHandler
			} else {
				urlinfo.addUrlPara(urlMacthPara{urlPara: para, matchInfo: URL_PARA_FLAG})
			}
		} else {
			matchParaCount--
			urlinfo.addUrlPara(urlMacthPara{urlPara: para, matchInfo: EMPTY})
		}
	}
	if matchParaCount == 0 {
		handlerMap.simple[url] = restHandler
	} else {
		urlinfo.handler = restHandler
		len := urlinfo.len()
		handlerMap.rest.urls[len] = append(handlerMap.rest.urls[len], urlinfo)
	}
}

/**
 * 添加一个过滤器
 */
func AddFilter(f Filter) {
	filters = append(filters, f)
}

/**
 * 实现http接口
 */
func (this *DoobHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	startTime := time.Now()
	defer log.Printf("程序处理共消耗:%d ns", time.Now().Sub(startTime).Nanoseconds())
	for i := range this.filters {
		if this.filters[i].Filter(res, req) {
			continue
		} else {
			return
		}
	}
	handler, err := this.handlerMap.getHandler(req)
	if err != nil {
		log.Println("error => ", err.Error())
		errStr := err.Error()
		if strings.Index(errStr, "method not match") >= 0 {
			res.WriteHeader(405)
		} else {
			res.WriteHeader(404)
		}
		return
	}
	handler(res, req)
}

func init() {
	simple := map[string]*restHandlerFunc{}
	rest := &restHandlerMap{urls: map[int][]*urlInfo{}}
	last := map[string]*restHandlerFunc{}
	handlerMap = &handleFuncMap{
		simple:       simple,
		rest:         rest,
		lastAllMatch: last,
	}
	filters = []Filter{}
	AddHandlerFunc("/", "", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("doob "))
	})
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/base64"
	"encoding/json"
	"strings"

	"github.com/astaxie/beego"
	"github.com/lfq7413/tomato/client"
	"github.com/lfq7413/tomato/config"
	"github.com/lfq7413/tomato/errs"
	"github.com/lfq7413/tomato/rest"
	"github.com/lfq7413/tomato/types"
	"github.com/lfq7413/tomato/utils"
)

// BaseController ...
// Info 当前请求的权限信息
// Auth 当前请求的用户权限
// JSONBody 由 JSON 格式转换来的请求数据
// RawBody 原始请求数据
type BaseController struct {
	beego.Controller
	Info     *RequestInfo
	Auth     *rest.Auth
	Query    map[string]string
	JSONBody types.M
	RawBody  []byte
}

// RequestInfo http 请求的权限信息
type RequestInfo struct {
	AppID          string
	MasterKey      string
	ClientKey      string
	JavascriptKey  string
	DotNetKey      string
	RestAPIKey     string
	SessionToken   string
	InstallationID string
	ClientVersion  string
	ClientSDK      map[string]string
}

// Prepare 对请求权限进行处理
// 1. 从请求头中获取各种 key
// 2. 尝试按 json 格式转换 body
// 3. 尝试从 body 中获取各种 key
// 4. 校验请求权限
// 5. 生成用户信息
func (b *BaseController) Prepare() {
	info := &RequestInfo{}
	info.AppID = b.Ctx.Input.Header("X-Parse-Application-Id")
	info.MasterKey = b.Ctx.Input.Header("X-Parse-Master-Key")
	info.ClientKey = b.Ctx.Input.Header("X-Parse-Client-Key")
	info.JavascriptKey = b.Ctx.Input.Header("X-Parse-Javascript-Key")
	info.DotNetKey = b.Ctx.Input.Header("X-Parse-Windows-Key")
	info.RestAPIKey = b.Ctx.Input.Header("X-Parse-REST-API-Key")
	info.SessionToken = b.Ctx.Input.Header("X-Parse-Session-Token")
	info.InstallationID = b.Ctx.Input.Header("X-Parse-Installation-Id")
	info.ClientVersion = b.Ctx.Input.Header("X-Parse-Client-Version")

	basicAuth := httpAuth(b.Ctx.Input.Header("Authorization"))
	if basicAuth != nil {
		info.AppID = basicAuth["appId"]
		if basicAuth["masterKey"] != "" {
			info.MasterKey = basicAuth["masterKey"]
		}
		if basicAuth["javascriptKey"] != "" {
			info.ClientKey = basicAuth["javascriptKey"]
		}
	}

	b.Query = map[string]string{}
	input := b.Input()
	for key := range input {
		b.Query[key] = input.Get(key)
	}

	if b.Ctx.Input.RequestBody != nil {
		contentType := b.Ctx.Input.Header("Content-type")
		if strings.HasPrefix(contentType, "application/json") {
			// 请求数据为 json 格式，进行转换，转换出错则返回错误
			var object types.M
			err := json.Unmarshal(b.Ctx.Input.RequestBody, &object)
			if err != nil {
				b.HandleError(errs.E(errs.InvalidJSON, "invalid JSON"), 0)
				return
			}
			b.JSONBody = object
		} else {
			// 当 AppID 不存在时，尝试转换，转换失败不返回错误
			if info.AppID == "" {
				var object types.M
				err := json.Unmarshal(b.Ctx.Input.RequestBody, &object)
				if err == nil {
					b.JSONBody = object
				}
			}
		}
	}

	if b.JSONBody != nil {
		// Unity SDK sends a _noBody key which needs to be removed.
		// Unclear at this point if action needs to be taken.
		delete(b.JSONBody, "_noBody")
		delete(b.JSONBody, "_method")
	}

	if info.AppID == "" {
		if b.JSONBody != nil {
			delete(b.JSONBody, "_RevocableSession")
		}
		// 从请求数据中获取各种 key
		if b.JSONBody != nil && b.JSONBody["_ApplicationId"] != nil {
			info.AppID = utils.S(b.JSONBody["_ApplicationId"])
			info.JavascriptKey = utils.S(b.JSONBody["_JavaScriptKey"])
			delete(b.JSONBody, "_ApplicationId")
			delete(b.JSONBody, "_JavaScriptKey")

			if b.JSONBody["_ClientVersion"] != nil {
				info.ClientVersion = utils.S(b.JSONBody["_ClientVersion"])
				delete(b.JSONBody, "_ClientVersion")
			}
			if b.JSONBody["_InstallationId"] != nil {
				info.InstallationID = utils.S(b.JSONBody["_InstallationId"])
				delete(b.JSONBody, "_InstallationId")
			}
			if b.JSONBody["_SessionToken"] != nil {
				info.SessionToken = utils.S(b.JSONBody["_SessionToken"])
				delete(b.JSONBody, "_SessionToken")
			}
			if b.JSONBody["_MasterKey"] != nil {
				info.MasterKey = utils.S(b.JSONBody["_MasterKey"])
				delete(b.JSONBody, "_MasterKey")
			}
			if b.JSONBody["_ContentType"] != nil {
				b.Ctx.Input.Context.Request.Header.Set("Content-type", utils.S(b.JSONBody["_ContentType"]))
				delete(b.JSONBody, "_ContentType")
			}
		} else {
			// 请求数据中也不存在 APPID 时，返回错误
			b.InvalidRequest()
			return
		}
	}

	if info.ClientVersion != "" {
		info.ClientSDK = client.FromString(info.ClientVersion)
	}

	b.Info = info

	// 校验请求权限
	if info.AppID != config.TConfig.AppID {
		b.InvalidRequest()
		return
	}
	if info.MasterKey == config.TConfig.MasterKey {
		b.Auth = &rest.Auth{InstallationID: info.InstallationID, IsMaster: true}
		return
	}
	var allow = false
	if (len(info.ClientKey) > 0 && info.ClientKey == config.TConfig.ClientKey) ||
		(len(info.JavascriptKey) > 0 && info.JavascriptKey == config.TConfig.JavascriptKey) ||
		(len(info.RestAPIKey) > 0 && info.RestAPIKey == config.TConfig.RestAPIKey) ||
		(len(info.DotNetKey) > 0 && info.DotNetKey == config.TConfig.DotNetKey) {
		allow = true
	}
	if allow == false {
		b.InvalidRequest()
		return
	}
	// TODO 登录时删除 Token ，如何处理接口地址？
	url := b.Ctx.Input.URL()
	if url == "/v1/login" || url == "/v1/login/" {
		info.SessionToken = ""
	}
	// 生成当前会话用户权限信息
	if info.SessionToken == "" {
		b.Auth = &rest.Auth{InstallationID: info.InstallationID, IsMaster: false}
		return
	}
	var auth *rest.Auth
	var err error
	if (url == "/v1/upgradeToRevocableSession" || url == "/v1/upgradeToRevocableSession/") &&
		strings.Index(info.SessionToken, "r:") != 0 {
		auth, err = rest.GetAuthForLegacySessionToken(info.SessionToken, info.InstallationID)
	} else {
		auth, err = rest.GetAuthForSessionToken(info.SessionToken, info.InstallationID)
	}
	if err != nil {
		b.HandleError(err, 0)
		return
	}
	b.Auth = auth
}

func httpAuth(authorization string) map[string]string {
	if authorization == "" {
		return nil
	}

	var appID, masterKey, javascriptKey string
	authPrefix1 := "basic "
	authPrefix2 := "Basic "

	match := strings.HasPrefix(authorization, authPrefix1) || strings.HasPrefix(authorization, authPrefix2)
	if match {
		encodedAuth := authorization[len(authPrefix1):]
		credentials := strings.Split(decodeBase64(encodedAuth), ":")

		if len(credentials) == 2 {
			appID = credentials[0]
			key := credentials[1]
			jsKeyPrefix := "javascript-key="

			matchKey := strings.HasPrefix(key, jsKeyPrefix)
			if matchKey {
				javascriptKey = key[len(jsKeyPrefix):]
			} else {
				masterKey = key
			}
			return map[string]string{
				"appId":         appID,
				"masterKey":     masterKey,
				"javascriptKey": javascriptKey,
			}
		}
		return nil
	}

	return nil
}

func decodeBase64(str string) string {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return ""
	}
	return string(data)
}

// HandleError 返回错误信息，不指定 status 参数时，默认为 0
func (b *BaseController) HandleError(err error, status int) {
	code := errs.GetErrorCode(err)
	if code != 0 {
		var httpStatus int
		switch code {
		case errs.InternalServerError:
			httpStatus = 500
		case errs.ObjectNotFound:
			httpStatus = 404
		default:
			httpStatus = 400
		}

		b.Ctx.Output.SetStatus(httpStatus)
		b.Data["json"] = errs.ErrorToMap(err)
		b.ServeJSON()
		return
	}

	if status != 0 {
		b.Ctx.Output.SetStatus(status)
		b.Data["json"] = types.M{"error": err.Error()}
		b.ServeJSON()
		return
	}

	b.Ctx.Output.SetStatus(500)
	b.Data["json"] = errs.ErrorMessageToMap(errs.InternalServerError, "Internal server error: "+err.Error())
	b.ServeJSON()
}

// InvalidRequest 无效请求
func (b *BaseController) InvalidRequest() {
	b.Ctx.Output.SetStatus(403)
	b.Data["json"] = types.M{"error": "unauthorized"}
	b.ServeJSON()
}

// EnforceMasterKeyAccess 接口需要 Master 权限
// 返回 true 表示当前请求是 Master 权限
func (b *BaseController) EnforceMasterKeyAccess() bool {
	if b.Auth.IsMaster == false {
		b.Ctx.Output.SetStatus(403)
		b.Data["json"] = types.M{"error": "unauthorized: master key is required"}
		b.ServeJSON()
		return false
	}
	return true
}
<commit_msg>兼容 Android SDK<commit_after>package controllers

import (
	"encoding/base64"
	"encoding/json"
	"strings"

	"github.com/astaxie/beego"
	"github.com/lfq7413/tomato/client"
	"github.com/lfq7413/tomato/config"
	"github.com/lfq7413/tomato/errs"
	"github.com/lfq7413/tomato/rest"
	"github.com/lfq7413/tomato/types"
	"github.com/lfq7413/tomato/utils"
)

// BaseController ...
// Info 当前请求的权限信息
// Auth 当前请求的用户权限
// JSONBody 由 JSON 格式转换来的请求数据
// RawBody 原始请求数据
type BaseController struct {
	beego.Controller
	Info     *RequestInfo
	Auth     *rest.Auth
	Query    map[string]string
	JSONBody types.M
	RawBody  []byte
}

// RequestInfo http 请求的权限信息
type RequestInfo struct {
	AppID          string
	MasterKey      string
	ClientKey      string
	JavascriptKey  string
	DotNetKey      string
	RestAPIKey     string
	SessionToken   string
	InstallationID string
	ClientVersion  string
	ClientSDK      map[string]string
}

// Prepare 对请求权限进行处理
// 1. 从请求头中获取各种 key
// 2. 尝试按 json 格式转换 body
// 3. 尝试从 body 中获取各种 key
// 4. 校验请求权限
// 5. 生成用户信息
func (b *BaseController) Prepare() {
	info := &RequestInfo{}
	info.AppID = b.Ctx.Input.Header("X-Parse-Application-Id")
	info.MasterKey = b.Ctx.Input.Header("X-Parse-Master-Key")
	info.ClientKey = b.Ctx.Input.Header("X-Parse-Client-Key")
	info.JavascriptKey = b.Ctx.Input.Header("X-Parse-Javascript-Key")
	info.DotNetKey = b.Ctx.Input.Header("X-Parse-Windows-Key")
	info.RestAPIKey = b.Ctx.Input.Header("X-Parse-REST-API-Key")
	info.SessionToken = b.Ctx.Input.Header("X-Parse-Session-Token")
	info.InstallationID = b.Ctx.Input.Header("X-Parse-Installation-Id")
	info.ClientVersion = b.Ctx.Input.Header("X-Parse-Client-Version")

	basicAuth := httpAuth(b.Ctx.Input.Header("Authorization"))
	if basicAuth != nil {
		info.AppID = basicAuth["appId"]
		if basicAuth["masterKey"] != "" {
			info.MasterKey = basicAuth["masterKey"]
		}
		if basicAuth["javascriptKey"] != "" {
			info.ClientKey = basicAuth["javascriptKey"]
		}
	}

	b.Query = map[string]string{}
	input := b.Input()
	for key := range input {
		b.Query[key] = input.Get(key)
	}

	if b.Ctx.Input.RequestBody != nil {
		contentType := b.Ctx.Input.Header("Content-type")
		if strings.HasPrefix(contentType, "application/json") {
			// 请求数据为 json 格式，进行转换，转换出错则返回错误
			var object types.M
			err := json.Unmarshal(b.Ctx.Input.RequestBody, &object)
			if err != nil {
				b.HandleError(errs.E(errs.InvalidJSON, "invalid JSON"), 0)
				return
			}
			b.JSONBody = object
		} else {
			// 当 AppID 不存在时，尝试转换，转换失败不返回错误
			if info.AppID == "" {
				var object types.M
				err := json.Unmarshal(b.Ctx.Input.RequestBody, &object)
				if err == nil {
					b.JSONBody = object
				}
			}
		}
	}

	if b.JSONBody != nil {
		// Unity SDK sends a _noBody key which needs to be removed.
		// Unclear at this point if action needs to be taken.
		delete(b.JSONBody, "_noBody")
		delete(b.JSONBody, "_method")
	}

	if info.AppID == "" {
		if b.JSONBody != nil {
			delete(b.JSONBody, "_RevocableSession")
		}
		// 从请求数据中获取各种 key
		if b.JSONBody != nil && b.JSONBody["_ApplicationId"] != nil {
			info.AppID = utils.S(b.JSONBody["_ApplicationId"])
			info.JavascriptKey = utils.S(b.JSONBody["_JavaScriptKey"])
			delete(b.JSONBody, "_ApplicationId")
			delete(b.JSONBody, "_JavaScriptKey")

			if b.JSONBody["_ClientVersion"] != nil {
				info.ClientVersion = utils.S(b.JSONBody["_ClientVersion"])
				delete(b.JSONBody, "_ClientVersion")
			}
			if b.JSONBody["_InstallationId"] != nil {
				info.InstallationID = utils.S(b.JSONBody["_InstallationId"])
				delete(b.JSONBody, "_InstallationId")
			}
			if b.JSONBody["_SessionToken"] != nil {
				info.SessionToken = utils.S(b.JSONBody["_SessionToken"])
				delete(b.JSONBody, "_SessionToken")
			}
			if b.JSONBody["_MasterKey"] != nil {
				info.MasterKey = utils.S(b.JSONBody["_MasterKey"])
				delete(b.JSONBody, "_MasterKey")
			}
			if b.JSONBody["_ContentType"] != nil {
				b.Ctx.Input.Context.Request.Header.Set("Content-type", utils.S(b.JSONBody["_ContentType"]))
				delete(b.JSONBody, "_ContentType")
			}
		} else {
			// 请求数据中也不存在 APPID 时，返回错误
			b.InvalidRequest()
			return
		}
	}

	// 兼容 Android SDK ，查询参数保存在 body 中，并且值的类型均为 string
	if b.Ctx.Input.Method() == "GET" {
		for key, value := range b.JSONBody {
			if str, ok := value.(string); ok {
				b.Query[key] = str
				delete(b.JSONBody, key)
			}
		}
	}

	if info.ClientVersion != "" {
		info.ClientSDK = client.FromString(info.ClientVersion)
	}

	b.Info = info

	// 校验请求权限
	if info.AppID != config.TConfig.AppID {
		b.InvalidRequest()
		return
	}
	if info.MasterKey == config.TConfig.MasterKey {
		b.Auth = &rest.Auth{InstallationID: info.InstallationID, IsMaster: true}
		return
	}
	var allow = false
	if (len(info.ClientKey) > 0 && info.ClientKey == config.TConfig.ClientKey) ||
		(len(info.JavascriptKey) > 0 && info.JavascriptKey == config.TConfig.JavascriptKey) ||
		(len(info.RestAPIKey) > 0 && info.RestAPIKey == config.TConfig.RestAPIKey) ||
		(len(info.DotNetKey) > 0 && info.DotNetKey == config.TConfig.DotNetKey) {
		allow = true
	}
	if allow == false {
		b.InvalidRequest()
		return
	}
	// TODO 登录时删除 Token ，如何处理接口地址？
	url := b.Ctx.Input.URL()
	if url == "/v1/login" || url == "/v1/login/" {
		info.SessionToken = ""
	}
	// 生成当前会话用户权限信息
	if info.SessionToken == "" {
		b.Auth = &rest.Auth{InstallationID: info.InstallationID, IsMaster: false}
		return
	}
	var auth *rest.Auth
	var err error
	if (url == "/v1/upgradeToRevocableSession" || url == "/v1/upgradeToRevocableSession/") &&
		strings.Index(info.SessionToken, "r:") != 0 {
		auth, err = rest.GetAuthForLegacySessionToken(info.SessionToken, info.InstallationID)
	} else {
		auth, err = rest.GetAuthForSessionToken(info.SessionToken, info.InstallationID)
	}
	if err != nil {
		b.HandleError(err, 0)
		return
	}
	b.Auth = auth
}

func httpAuth(authorization string) map[string]string {
	if authorization == "" {
		return nil
	}

	var appID, masterKey, javascriptKey string
	authPrefix1 := "basic "
	authPrefix2 := "Basic "

	match := strings.HasPrefix(authorization, authPrefix1) || strings.HasPrefix(authorization, authPrefix2)
	if match {
		encodedAuth := authorization[len(authPrefix1):]
		credentials := strings.Split(decodeBase64(encodedAuth), ":")

		if len(credentials) == 2 {
			appID = credentials[0]
			key := credentials[1]
			jsKeyPrefix := "javascript-key="

			matchKey := strings.HasPrefix(key, jsKeyPrefix)
			if matchKey {
				javascriptKey = key[len(jsKeyPrefix):]
			} else {
				masterKey = key
			}
			return map[string]string{
				"appId":         appID,
				"masterKey":     masterKey,
				"javascriptKey": javascriptKey,
			}
		}
		return nil
	}

	return nil
}

func decodeBase64(str string) string {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return ""
	}
	return string(data)
}

// HandleError 返回错误信息，不指定 status 参数时，默认为 0
func (b *BaseController) HandleError(err error, status int) {
	code := errs.GetErrorCode(err)
	if code != 0 {
		var httpStatus int
		switch code {
		case errs.InternalServerError:
			httpStatus = 500
		case errs.ObjectNotFound:
			httpStatus = 404
		default:
			httpStatus = 400
		}

		b.Ctx.Output.SetStatus(httpStatus)
		b.Data["json"] = errs.ErrorToMap(err)
		b.ServeJSON()
		return
	}

	if status != 0 {
		b.Ctx.Output.SetStatus(status)
		b.Data["json"] = types.M{"error": err.Error()}
		b.ServeJSON()
		return
	}

	b.Ctx.Output.SetStatus(500)
	b.Data["json"] = errs.ErrorMessageToMap(errs.InternalServerError, "Internal server error: "+err.Error())
	b.ServeJSON()
}

// InvalidRequest 无效请求
func (b *BaseController) InvalidRequest() {
	b.Ctx.Output.SetStatus(403)
	b.Data["json"] = types.M{"error": "unauthorized"}
	b.ServeJSON()
}

// EnforceMasterKeyAccess 接口需要 Master 权限
// 返回 true 表示当前请求是 Master 权限
func (b *BaseController) EnforceMasterKeyAccess() bool {
	if b.Auth.IsMaster == false {
		b.Ctx.Output.SetStatus(403)
		b.Data["json"] = types.M{"error": "unauthorized: master key is required"}
		b.ServeJSON()
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package controllers

import (
	"github.com/codegangsta/martini-contrib/render"
	"github.com/martini-contrib/oauth2"
	"net/http"
	"io/ioutil"
)

// Get the logged in user
func GetLoggedInUser(tokens oauth2.Tokens, r render.Render) {

	var url = "https://www.googleapis.com/plus/v1/people/me?access_token=" + tokens.Access()

	resp, err := http.Get(url)
	if err != nil {
		panic(err)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}

	return body

}
<commit_msg>Add a response code<commit_after>package controllers

import (
	"github.com/codegangsta/martini-contrib/render"
	"github.com/martini-contrib/oauth2"
	"net/http"
	"io/ioutil"
)

// Get the logged in user
func GetLoggedInUser(tokens oauth2.Tokens, r render.Render) (int,string) {

	var url = "https://www.googleapis.com/plus/v1/people/me?access_token=" + tokens.Access()

	resp, err := http.Get(url)
	if err != nil {
		panic(err)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}

	return 200, string(body)
}
<|endoftext|>


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

package controllers

import (
	"github.com/Unknwon/gcblog/models"
)

type WorkController struct {
	baseController
}

func (this *WorkController) Get() {
	this.Data["IsWork"] = true
	this.TplNames = "home.html"
	this.Data["AllWorks"] = models.GetAllWorks()
	this.Data["RecentArchives"] = models.GetAllPosts()
}

func (this *WorkController) SingleWork() {
	this.TplNames = "home.html"
	this.Data["IsWork"] = true
	this.Data["IsSingleWork"] = true
	this.Data["RecentArchives"] = models.GetAllPosts()
	work := models.GetSingleWork(this.Ctx.Request.RequestURI[6:])
	if work == nil {
		this.Redirect("/", 302)
		return
	}
	this.Data["SingleWork"] = work
}
<commit_msg>Update code<commit_after>// Copyright 2013 gcblog authors.
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package controllers

import (
	"github.com/Unknwon/gcblog/models"
)

type WorkController struct {
	baseController
}

func (this *WorkController) Get() {
	this.Data["IsWork"] = true
	this.TplNames = "home.html"
	this.Data["AllWorks"] = models.GetAllWorks()
	this.Data["RecentList"] = models.GetRecentPosts(10)
}

func (this *WorkController) SingleWork() {
	this.TplNames = "home.html"
	this.Data["IsWork"] = true
	this.Data["IsSingleWork"] = true
	this.Data["RecentList"] = models.GetRecentPosts(10)
	work := models.GetSingleWork(this.Ctx.Request.RequestURI[6:])
	if work == nil {
		this.Redirect("/", 302)
		return
	}
	this.Data["SingleWork"] = work
}
<|endoftext|>


<commit_before>package copy

import (
	"io"

	internalblobinfocache "github.com/containers/image/v5/internal/blobinfocache"
	"github.com/containers/image/v5/pkg/compression"
	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/image/v5/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// bpDetectCompressionStepData contains data that the copy pipeline needs about the “detect compression” step.
type bpDetectCompressionStepData struct {
	isCompressed bool
	format       compressiontypes.Algorithm        // Valid if isCompressed
	decompressor compressiontypes.DecompressorFunc // Valid if isCompressed
}

// blobPipelineDetectCompressionStep updates *stream to detect its current compression format.
// srcInfo is only used for error messages.
// Returns data for other steps.
func blobPipelineDetectCompressionStep(stream *sourceStream, srcInfo types.BlobInfo) (bpDetectCompressionStepData, error) {
	// This requires us to “peek ahead” into the stream to read the initial part, which requires us to chain through another io.Reader returned by DetectCompression.
	format, decompressor, reader, err := compression.DetectCompressionFormat(stream.reader) // We could skip this in some cases, but let's keep the code path uniform
	if err != nil {
		return bpDetectCompressionStepData{}, errors.Wrapf(err, "reading blob %s", srcInfo.Digest)
	}
	stream.reader = reader

	res := bpDetectCompressionStepData{
		isCompressed: decompressor != nil,
		format:       format,
		decompressor: decompressor,
	}
	if expectedFormat, known := expectedCompressionFormats[stream.info.MediaType]; known && res.isCompressed && format.Name() != expectedFormat.Name() {
		logrus.Debugf("blob %s with type %s should be compressed with %s, but compressor appears to be %s", srcInfo.Digest.String(), srcInfo.MediaType, expectedFormat.Name(), format.Name())
	}
	return res, nil
}

// bpCompressionStepData contains data that the copy pipeline needs about the compression step.
type bpCompressionStepData struct {
	operation              types.LayerCompression      // Operation to use for updating the blob metadata.
	uploadedAlgorithm      *compressiontypes.Algorithm // An algorithm parameter for the compressionOperation edits.
	uploadedAnnotations    map[string]string           // Annotations that should be set on the uploaded blob. WARNING: This is only set after the srcStream.reader is fully consumed.
	srcCompressorName      string                      // Compressor name to record in the blob info cache for the source blob.
	uploadedCompressorName string                      // Compressor name to record in the blob info cache for the uploaded blob.
	closers                []io.Closer                 // Objects to close after the upload is done, if any.
}

// blobPipelineCompressionStep updates *stream to compress and/or decompress it.
// srcInfo is only used for error messages.
// Returns data for other steps; the caller should eventually call updateCompressionEdits and perhaps recordValidatedBlobData,
// and must eventually call close.
func (c *copier) blobPipelineCompressionStep(stream *sourceStream, canModifyBlob bool,
	detected bpDetectCompressionStepData) (*bpCompressionStepData, error) {
	// WARNING: If you are adding new reasons to change the blob, update also the OptimizeDestinationImageAlreadyExists
	// short-circuit conditions
	uploadedAnnotations := map[string]string{}
	var operation types.LayerCompression
	var uploadedAlgorithm *compressiontypes.Algorithm
	srcCompressorName := internalblobinfocache.Uncompressed
	if detected.isCompressed {
		srcCompressorName = detected.format.Name()
	}
	var uploadedCompressorName string
	var closers []io.Closer
	succeeded := false
	defer func() {
		if !succeeded {
			for _, c := range closers {
				c.Close()
			}
		}
	}()
	if canModifyBlob && isOciEncrypted(stream.info.MediaType) {
		// PreserveOriginal due to any compression not being able to be done on an encrypted blob unless decrypted
		logrus.Debugf("Using original blob without modification for encrypted blob")
		operation = types.PreserveOriginal
		srcCompressorName = internalblobinfocache.UnknownCompression
		uploadedAlgorithm = nil
		uploadedCompressorName = internalblobinfocache.UnknownCompression
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && !detected.isCompressed {
		logrus.Debugf("Compressing blob on the fly")
		operation = types.Compress
		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		if c.compressionFormat != nil {
			uploadedAlgorithm = c.compressionFormat
		} else {
			uploadedAlgorithm = defaultCompressionFormat
		}
		// If this fails while writing data, it will do pipeWriter.CloseWithError(); if it fails otherwise,
		// e.g. because we have exited and due to pipeReader.Close() above further writing to the pipe has failed,
		// we don’t care.
		go c.compressGoroutine(pipeWriter, stream.reader, uploadedAnnotations, *uploadedAlgorithm) // Closes pipeWriter
		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && detected.isCompressed &&
		c.compressionFormat != nil && c.compressionFormat.Name() != detected.format.Name() {
		// When the blob is compressed, but the desired format is different, it first needs to be decompressed and finally
		// re-compressed using the desired format.
		logrus.Debugf("Blob will be converted")

		operation = types.PreserveOriginal
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)

		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		uploadedAlgorithm = c.compressionFormat
		go c.compressGoroutine(pipeWriter, s, uploadedAnnotations, *uploadedAlgorithm) // Closes pipeWriter

		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Decompress && detected.isCompressed {
		logrus.Debugf("Blob will be decompressed")
		operation = types.Decompress
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)
		stream.reader = s
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedAlgorithm = nil
		uploadedCompressorName = internalblobinfocache.Uncompressed
	} else {
		// PreserveOriginal might also need to recompress the original blob if the desired compression format is different.
		logrus.Debugf("Using original blob without modification")
		operation = types.PreserveOriginal
		// Remember if the original blob was compressed, and if so how, so that if
		// LayerInfosForCopy() returned something that differs from what was in the
		// source's manifest, and UpdatedImage() needs to call UpdateLayerInfos(),
		// it will be able to correctly derive the MediaType for the copied blob.
		if detected.isCompressed {
			uploadedAlgorithm = &detected.format
		} else {
			uploadedAlgorithm = nil
		}
		uploadedCompressorName = srcCompressorName
	}
	succeeded = true
	return &bpCompressionStepData{
		operation:              operation,
		uploadedAlgorithm:      uploadedAlgorithm,
		uploadedAnnotations:    uploadedAnnotations,
		srcCompressorName:      srcCompressorName,
		uploadedCompressorName: uploadedCompressorName,
		closers:                closers,
	}, nil
}

// updateCompressionEdits sets *operation, *algorithm and updates *annotations, if necessary.
func (d *bpCompressionStepData) updateCompressionEdits(operation *types.LayerCompression, algorithm **compressiontypes.Algorithm, annotations *map[string]string) {
	*operation = d.operation
	// If we can modify the layer's blob, set the desired algorithm for it to be set in the manifest.
	*algorithm = d.uploadedAlgorithm
	if *annotations == nil {
		*annotations = map[string]string{}
	}
	for k, v := range d.uploadedAnnotations {
		(*annotations)[k] = v
	}
}

// recordValidatedBlobData updates b.blobInfoCache with data about the created uploadedInfo adnd the original srcInfo.
// This must ONLY be called if all data has been validated by OUR code, and is not comming from third parties.
func (d *bpCompressionStepData) recordValidatedDigestData(c *copier, uploadedInfo types.BlobInfo, srcInfo types.BlobInfo,
	encryptionStep *bpEncryptionStepData, decryptionStep *bpDecryptionStepData) error {
	// Don’t record any associations that involve encrypted data. This is a bit crude,
	// some blob substitutions (replacing pulls of encrypted data with local reuse of known decryption outcomes)
	// might be safe, but it’s not trivially obvious, so let’s be conservative for now.
	// This crude approach also means we don’t need to record whether a blob is encrypted
	// in the blob info cache (which would probably be necessary for any more complex logic),
	// and the simplicity is attractive.
	if !encryptionStep.encrypting && !decryptionStep.decrypting {
		// If d.operation != types.PreserveOriginal, we now have two reliable digest values:
		// srcinfo.Digest describes the pre-d.operation input, verified by digestingReader
		// uploadedInfo.Digest describes the post-d.operation output, computed by PutBlob
		// (because stream.info.Digest == "", this must have been computed afresh).
		switch d.operation {
		case types.PreserveOriginal:
			break // Do nothing, we have only one digest and we might not have even verified it.
		case types.Compress:
			c.blobInfoCache.RecordDigestUncompressedPair(uploadedInfo.Digest, srcInfo.Digest)
		case types.Decompress:
			c.blobInfoCache.RecordDigestUncompressedPair(srcInfo.Digest, uploadedInfo.Digest)
		default:
			return errors.Errorf("Internal error: Unexpected d.operation value %#v", d.operation)
		}
	}
	if d.uploadedCompressorName != "" && d.uploadedCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(uploadedInfo.Digest, d.uploadedCompressorName)
	}
	if srcInfo.Digest != "" && d.srcCompressorName != "" && d.srcCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(srcInfo.Digest, d.srcCompressorName)
	}
	return nil
}

// close closes objects that carry state throughout the compression/decompression operation.
func (d *bpCompressionStepData) close() {
	for _, c := range d.closers {
		c.Close()
	}
}

// doCompression reads all input from src and writes its compressed equivalent to dest.
func doCompression(dest io.Writer, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm, compressionLevel *int) error {
	compressor, err := compression.CompressStreamWithMetadata(dest, metadata, compressionFormat, compressionLevel)
	if err != nil {
		return err
	}

	buf := make([]byte, compressionBufferSize)

	_, err = io.CopyBuffer(compressor, src, buf) // Sets err to nil, i.e. causes dest.Close()
	if err != nil {
		compressor.Close()
		return err
	}

	return compressor.Close()
}

// compressGoroutine reads all input from src and writes its compressed equivalent to dest.
func (c *copier) compressGoroutine(dest *io.PipeWriter, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm) {
	err := errors.New("Internal error: unexpected panic in compressGoroutine")
	defer func() { // Note that this is not the same as {defer dest.CloseWithError(err)}; we need err to be evaluated lazily.
		_ = dest.CloseWithError(err) // CloseWithError(nil) is equivalent to Close(), always returns nil
	}()

	err = doCompression(dest, src, metadata, compressionFormat, c.compressionLevel)
}
<commit_msg>Compute srcCompressorName already in blobPipelineDetectCompressionStep<commit_after>package copy

import (
	"io"

	internalblobinfocache "github.com/containers/image/v5/internal/blobinfocache"
	"github.com/containers/image/v5/pkg/compression"
	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/image/v5/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// bpDetectCompressionStepData contains data that the copy pipeline needs about the “detect compression” step.
type bpDetectCompressionStepData struct {
	isCompressed      bool
	format            compressiontypes.Algorithm        // Valid if isCompressed
	decompressor      compressiontypes.DecompressorFunc // Valid if isCompressed
	srcCompressorName string                            // Compressor name to possibly record in the blob info cache for the source blob.
}

// blobPipelineDetectCompressionStep updates *stream to detect its current compression format.
// srcInfo is only used for error messages.
// Returns data for other steps.
func blobPipelineDetectCompressionStep(stream *sourceStream, srcInfo types.BlobInfo) (bpDetectCompressionStepData, error) {
	// This requires us to “peek ahead” into the stream to read the initial part, which requires us to chain through another io.Reader returned by DetectCompression.
	format, decompressor, reader, err := compression.DetectCompressionFormat(stream.reader) // We could skip this in some cases, but let's keep the code path uniform
	if err != nil {
		return bpDetectCompressionStepData{}, errors.Wrapf(err, "reading blob %s", srcInfo.Digest)
	}
	stream.reader = reader

	res := bpDetectCompressionStepData{
		isCompressed: decompressor != nil,
		format:       format,
		decompressor: decompressor,
	}
	if res.isCompressed {
		res.srcCompressorName = format.Name()
	} else {
		res.srcCompressorName = internalblobinfocache.Uncompressed
	}

	if expectedFormat, known := expectedCompressionFormats[stream.info.MediaType]; known && res.isCompressed && format.Name() != expectedFormat.Name() {
		logrus.Debugf("blob %s with type %s should be compressed with %s, but compressor appears to be %s", srcInfo.Digest.String(), srcInfo.MediaType, expectedFormat.Name(), format.Name())
	}
	return res, nil
}

// bpCompressionStepData contains data that the copy pipeline needs about the compression step.
type bpCompressionStepData struct {
	operation              types.LayerCompression      // Operation to use for updating the blob metadata.
	uploadedAlgorithm      *compressiontypes.Algorithm // An algorithm parameter for the compressionOperation edits.
	uploadedAnnotations    map[string]string           // Annotations that should be set on the uploaded blob. WARNING: This is only set after the srcStream.reader is fully consumed.
	srcCompressorName      string                      // Compressor name to record in the blob info cache for the source blob.
	uploadedCompressorName string                      // Compressor name to record in the blob info cache for the uploaded blob.
	closers                []io.Closer                 // Objects to close after the upload is done, if any.
}

// blobPipelineCompressionStep updates *stream to compress and/or decompress it.
// srcInfo is only used for error messages.
// Returns data for other steps; the caller should eventually call updateCompressionEdits and perhaps recordValidatedBlobData,
// and must eventually call close.
func (c *copier) blobPipelineCompressionStep(stream *sourceStream, canModifyBlob bool,
	detected bpDetectCompressionStepData) (*bpCompressionStepData, error) {
	// WARNING: If you are adding new reasons to change the blob, update also the OptimizeDestinationImageAlreadyExists
	// short-circuit conditions
	uploadedAnnotations := map[string]string{}
	var operation types.LayerCompression
	var uploadedAlgorithm *compressiontypes.Algorithm
	srcCompressorName := detected.srcCompressorName
	var uploadedCompressorName string
	var closers []io.Closer
	succeeded := false
	defer func() {
		if !succeeded {
			for _, c := range closers {
				c.Close()
			}
		}
	}()
	if canModifyBlob && isOciEncrypted(stream.info.MediaType) {
		// PreserveOriginal due to any compression not being able to be done on an encrypted blob unless decrypted
		logrus.Debugf("Using original blob without modification for encrypted blob")
		operation = types.PreserveOriginal
		srcCompressorName = internalblobinfocache.UnknownCompression
		uploadedAlgorithm = nil
		uploadedCompressorName = internalblobinfocache.UnknownCompression
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && !detected.isCompressed {
		logrus.Debugf("Compressing blob on the fly")
		operation = types.Compress
		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		if c.compressionFormat != nil {
			uploadedAlgorithm = c.compressionFormat
		} else {
			uploadedAlgorithm = defaultCompressionFormat
		}
		// If this fails while writing data, it will do pipeWriter.CloseWithError(); if it fails otherwise,
		// e.g. because we have exited and due to pipeReader.Close() above further writing to the pipe has failed,
		// we don’t care.
		go c.compressGoroutine(pipeWriter, stream.reader, uploadedAnnotations, *uploadedAlgorithm) // Closes pipeWriter
		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Compress && detected.isCompressed &&
		c.compressionFormat != nil && c.compressionFormat.Name() != detected.format.Name() {
		// When the blob is compressed, but the desired format is different, it first needs to be decompressed and finally
		// re-compressed using the desired format.
		logrus.Debugf("Blob will be converted")

		operation = types.PreserveOriginal
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)

		pipeReader, pipeWriter := io.Pipe()
		closers = append(closers, pipeReader)

		uploadedAlgorithm = c.compressionFormat
		go c.compressGoroutine(pipeWriter, s, uploadedAnnotations, *uploadedAlgorithm) // Closes pipeWriter

		stream.reader = pipeReader
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedCompressorName = uploadedAlgorithm.Name()
	} else if canModifyBlob && c.dest.DesiredLayerCompression() == types.Decompress && detected.isCompressed {
		logrus.Debugf("Blob will be decompressed")
		operation = types.Decompress
		s, err := detected.decompressor(stream.reader)
		if err != nil {
			return nil, err
		}
		closers = append(closers, s)
		stream.reader = s
		stream.info = types.BlobInfo{ // FIXME? Should we preserve more data in src.info?
			Digest: "",
			Size:   -1,
		}
		uploadedAlgorithm = nil
		uploadedCompressorName = internalblobinfocache.Uncompressed
	} else {
		// PreserveOriginal might also need to recompress the original blob if the desired compression format is different.
		logrus.Debugf("Using original blob without modification")
		operation = types.PreserveOriginal
		// Remember if the original blob was compressed, and if so how, so that if
		// LayerInfosForCopy() returned something that differs from what was in the
		// source's manifest, and UpdatedImage() needs to call UpdateLayerInfos(),
		// it will be able to correctly derive the MediaType for the copied blob.
		if detected.isCompressed {
			uploadedAlgorithm = &detected.format
		} else {
			uploadedAlgorithm = nil
		}
		uploadedCompressorName = srcCompressorName
	}
	succeeded = true
	return &bpCompressionStepData{
		operation:              operation,
		uploadedAlgorithm:      uploadedAlgorithm,
		uploadedAnnotations:    uploadedAnnotations,
		srcCompressorName:      srcCompressorName,
		uploadedCompressorName: uploadedCompressorName,
		closers:                closers,
	}, nil
}

// updateCompressionEdits sets *operation, *algorithm and updates *annotations, if necessary.
func (d *bpCompressionStepData) updateCompressionEdits(operation *types.LayerCompression, algorithm **compressiontypes.Algorithm, annotations *map[string]string) {
	*operation = d.operation
	// If we can modify the layer's blob, set the desired algorithm for it to be set in the manifest.
	*algorithm = d.uploadedAlgorithm
	if *annotations == nil {
		*annotations = map[string]string{}
	}
	for k, v := range d.uploadedAnnotations {
		(*annotations)[k] = v
	}
}

// recordValidatedBlobData updates b.blobInfoCache with data about the created uploadedInfo adnd the original srcInfo.
// This must ONLY be called if all data has been validated by OUR code, and is not comming from third parties.
func (d *bpCompressionStepData) recordValidatedDigestData(c *copier, uploadedInfo types.BlobInfo, srcInfo types.BlobInfo,
	encryptionStep *bpEncryptionStepData, decryptionStep *bpDecryptionStepData) error {
	// Don’t record any associations that involve encrypted data. This is a bit crude,
	// some blob substitutions (replacing pulls of encrypted data with local reuse of known decryption outcomes)
	// might be safe, but it’s not trivially obvious, so let’s be conservative for now.
	// This crude approach also means we don’t need to record whether a blob is encrypted
	// in the blob info cache (which would probably be necessary for any more complex logic),
	// and the simplicity is attractive.
	if !encryptionStep.encrypting && !decryptionStep.decrypting {
		// If d.operation != types.PreserveOriginal, we now have two reliable digest values:
		// srcinfo.Digest describes the pre-d.operation input, verified by digestingReader
		// uploadedInfo.Digest describes the post-d.operation output, computed by PutBlob
		// (because stream.info.Digest == "", this must have been computed afresh).
		switch d.operation {
		case types.PreserveOriginal:
			break // Do nothing, we have only one digest and we might not have even verified it.
		case types.Compress:
			c.blobInfoCache.RecordDigestUncompressedPair(uploadedInfo.Digest, srcInfo.Digest)
		case types.Decompress:
			c.blobInfoCache.RecordDigestUncompressedPair(srcInfo.Digest, uploadedInfo.Digest)
		default:
			return errors.Errorf("Internal error: Unexpected d.operation value %#v", d.operation)
		}
	}
	if d.uploadedCompressorName != "" && d.uploadedCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(uploadedInfo.Digest, d.uploadedCompressorName)
	}
	if srcInfo.Digest != "" && d.srcCompressorName != "" && d.srcCompressorName != internalblobinfocache.UnknownCompression {
		c.blobInfoCache.RecordDigestCompressorName(srcInfo.Digest, d.srcCompressorName)
	}
	return nil
}

// close closes objects that carry state throughout the compression/decompression operation.
func (d *bpCompressionStepData) close() {
	for _, c := range d.closers {
		c.Close()
	}
}

// doCompression reads all input from src and writes its compressed equivalent to dest.
func doCompression(dest io.Writer, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm, compressionLevel *int) error {
	compressor, err := compression.CompressStreamWithMetadata(dest, metadata, compressionFormat, compressionLevel)
	if err != nil {
		return err
	}

	buf := make([]byte, compressionBufferSize)

	_, err = io.CopyBuffer(compressor, src, buf) // Sets err to nil, i.e. causes dest.Close()
	if err != nil {
		compressor.Close()
		return err
	}

	return compressor.Close()
}

// compressGoroutine reads all input from src and writes its compressed equivalent to dest.
func (c *copier) compressGoroutine(dest *io.PipeWriter, src io.Reader, metadata map[string]string, compressionFormat compressiontypes.Algorithm) {
	err := errors.New("Internal error: unexpected panic in compressGoroutine")
	defer func() { // Note that this is not the same as {defer dest.CloseWithError(err)}; we need err to be evaluated lazily.
		_ = dest.CloseWithError(err) // CloseWithError(nil) is equivalent to Close(), always returns nil
	}()

	err = doCompression(dest, src, metadata, compressionFormat, c.compressionLevel)
}
<|endoftext|>


<commit_before>package tests

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/ghodss/yaml"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/containerized-data-importer/tests/framework"
	"kubevirt.io/containerized-data-importer/tests/utils"
)

const (
	dataVolumeName     = "test-dv"
	pvcName            = "test-pvc"
	validURL           = "http://www.example.com/example.img"
	invalidURLFormat   = "invalidURL"
	datavolumeTestFile = "manifests/datavolume.yaml"
	destinationFile    = "/var/tmp/datavolume_test.yaml"
)

var _ = Describe("[rfe_id:1130][crit:medium][vendor:cnv-qe@redhat.com][level:component]Validation tests", func() {
	f := framework.NewFrameworkOrDie("api-validation-func-test")

	Describe("[posneg:negative]Verify DataVolume validation", func() {
		Context("when creating Datavolume", func() {
			dv := map[string]interface{}{}

			AfterEach(func() {
				err := os.Remove(destinationFile)
				Expect(err).ToNot(HaveOccurred())
			})

			table.DescribeTable("with Datavolume source validation should", func(sourceType string, args ...string) {

				By("Reading yaml file from: " + datavolumeTestFile)
				err := yamlFiletoStruct(datavolumeTestFile, &dv)
				Expect(err).ToNot(HaveOccurred())

				switch sourceType {
				case "http":
					url := args[0]
					dv["spec"].(map[string]interface{})["source"] = map[string]interface{}{"http": map[string]interface{}{"url": url}}

				case "s3":
					url := args[0]
					dv["spec"].(map[string]interface{})["source"] = map[string]interface{}{"s3": map[string]interface{}{"url": url}}
				case "pvc":
					namespace := args[0]
					name := args[1]
					dv["spec"].(map[string]interface{})["source"] = map[string]interface{}{
						"pvc": map[string]interface{}{
							"namespace": namespace,
							"name":      name}}
				}

				err = structToYamlFile(destinationFile, dv)
				Expect(err).ToNot(HaveOccurred())

				By(fmt.Sprint("Verifying kubectl create"))
				Eventually(func() bool {
					_, err := RunKubectlCommand(f, "create", "-f", destinationFile, "-n", f.Namespace.Name)
					if err != nil {
						return true
					}
					return false
				}, timeout, pollingInterval).Should(BeTrue())

			},
				table.Entry("[test_id:1321]fail with http source with invalid url format", "http", invalidURLFormat),
				table.Entry("[test_id:1322]fail with http source with empty url", "http", ""),
				table.Entry("[test_id:1323][crit:low]fail with s3 source with invalid url format", "s3", invalidURLFormat),
				table.Entry("[test_id:1324][crit:low]fail with s3 source with empty url", "s3", ""),
				table.Entry("[test_id:1325]fail with empty PVC source namespace", "pvc", "", "test-pvc"),
				table.Entry("[test_id:1326]fail with empty PVC source name", "pvc", "test", ""),
			)

			table.DescribeTable("with Datavolume PVC size should", func(size string) {

				By("Reading yaml file from: " + datavolumeTestFile)
				err := yamlFiletoStruct(datavolumeTestFile, &dv)
				Expect(err).ToNot(HaveOccurred())

				dv["spec"].(map[string]interface{})["pvc"].(map[string]interface{})["resources"].(map[string]interface{})["requests"].(map[string]interface{})["storage"] = size
				err = structToYamlFile(destinationFile, dv)
				Expect(err).ToNot(HaveOccurred())

				By(fmt.Sprint("Verifying kubectl apply"))
				Eventually(func() bool {
					_, err := RunKubectlCommand(f, "create", "-f", destinationFile, "-n", f.Namespace.Name)
					if err != nil {
						return true
					}
					return false
				}, timeout, pollingInterval).Should(BeTrue())
			},
				table.Entry("[test_id:1033]fail with zero PVC size", "0"),
				table.Entry("[test_id:1327]fail with negative PVC size", "-500m"),
				table.Entry("[test_id:1328]fail with invalid PVC size", "invalid_size"),
			)

		})
	})

	Context("DataVolume Already Exists", func() {
		BeforeEach(func() {
			dataVolume := utils.NewDataVolumeWithHTTPImport(dataVolumeName, "500Mi", validURL)

			dataVolume, err := utils.CreateDataVolumeFromDefinition(f.CdiClient, f.Namespace.Name, dataVolume)
			Expect(err).ToNot(HaveOccurred())
		})

		AfterEach(func() {
			err := utils.DeleteDataVolume(f.CdiClient, f.Namespace.Name, dataVolumeName)
			Expect(err).ToNot(HaveOccurred())
		})
		It("[test_id:1030]should fail creating an already existing DataVolume", func() {
			By(fmt.Sprint("Verifying kubectl create"))
			Eventually(func() bool {

				_, err := RunKubectlCommand(f, "create", "-f", datavolumeTestFile, "-n", f.Namespace.Name)
				if err != nil {
					return true
				}
				return false
			}, timeout, pollingInterval).Should(BeTrue())

		})
	})

	Context("DataVolume destination PVC", func() {
		BeforeEach(func() {
			pvc := utils.NewPVCDefinition(dataVolumeName, "50Mi", nil, nil)

			pvc, err := utils.CreatePVCFromDefinition(f.K8sClient, f.Namespace.Name, pvc)
			Expect(err).ToNot(HaveOccurred())
		})

		AfterEach(func() {
			pvc, err := f.K8sClient.CoreV1().PersistentVolumeClaims(f.Namespace.Name).Get(dataVolumeName, metav1.GetOptions{})
			Expect(err).ToNot(HaveOccurred())
			err = utils.DeletePVC(f.K8sClient, f.Namespace.Name, pvc)
			Expect(err).ToNot(HaveOccurred())
		})
		It("should fail creating a DataVolume with already existing destination pvc", func() {
			By(fmt.Sprint("Verifying kubectl create"))
			Eventually(func() bool {

				_, err := RunKubectlCommand(f, "create", "-f", datavolumeTestFile, "-n", f.Namespace.Name)
				if err != nil {
					return true
				}
				return false
			}, timeout, pollingInterval).Should(BeTrue())

		})
	})

	Context("when creating data volumes from manual manifests", func() {
		table.DescribeTable("with manifests Datavolume should", func(destinationFile string, expectError bool) {
			By(fmt.Sprint("Verifying kubectl apply"))
			_, err := RunKubectlCommand(f, "create", "-f", destinationFile, "-n", f.Namespace.Name)
			if expectError {
				Expect(err).To(HaveOccurred())
			} else {
				Expect(err).ToNot(HaveOccurred())
			}
		},
			table.Entry("fail with blank image source and contentType archive", "manifests/dvBlankArchive.yaml", true),
			table.Entry("fail with invalid contentType", "manifests/dvInvalidContentType.yaml", true),
			table.Entry("fail with missing source", "manifests/dvMissingSource.yaml", true),
			table.Entry("fail with multiple sources", "manifests/dvMultiSource.yaml", true),
			table.Entry("fail with invalid URL for http source", "manifests/dvInvalidURL.yaml", true),
			table.Entry("fail with invalid source PVC", "manifests/dvInvalidSourcePVC.yaml", true),
			table.Entry("succeed with valid source http", "manifests/datavolume.yaml", false),
			table.Entry("fail with missing PVC spec", "manifests/dvMissingPVCSpec.yaml", true),
			table.Entry("fail with missing resources spec", "manifests/dvMissingResourcesSpec.yaml", true),
			table.Entry("fail with 0 size PVC", "manifests/dv0SizePVC.yaml", true),
		)

	})
})

func yamlFiletoStruct(fileName string, o *map[string]interface{}) error {
	yamlFile, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}
	err = yaml.Unmarshal(yamlFile, o)
	if err != nil {
		return err
	}
	return nil

}

func structToYamlFile(fileName string, o interface{}) error {
	yamlOutput, err := yaml.Marshal(o)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(fileName, yamlOutput, 0644)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>add functional test for DataVolume validation using non existing source PVC<commit_after>package tests

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/ghodss/yaml"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/containerized-data-importer/tests/framework"
	"kubevirt.io/containerized-data-importer/tests/utils"
)

const (
	dataVolumeName     = "test-dv"
	pvcName            = "test-pvc"
	validURL           = "http://www.example.com/example.img"
	invalidURLFormat   = "invalidURL"
	datavolumeTestFile = "manifests/datavolume.yaml"
	destinationFile    = "/var/tmp/datavolume_test.yaml"
)

var _ = Describe("[rfe_id:1130][crit:medium][vendor:cnv-qe@redhat.com][level:component]Validation tests", func() {
	f := framework.NewFrameworkOrDie("api-validation-func-test")

	Describe("[posneg:negative]Verify DataVolume validation", func() {
		Context("when creating Datavolume", func() {
			dv := map[string]interface{}{}

			AfterEach(func() {
				err := os.Remove(destinationFile)
				Expect(err).ToNot(HaveOccurred())
			})

			table.DescribeTable("with Datavolume source validation should", func(sourceType string, args ...string) {

				By("Reading yaml file from: " + datavolumeTestFile)
				err := yamlFiletoStruct(datavolumeTestFile, &dv)
				Expect(err).ToNot(HaveOccurred())

				switch sourceType {
				case "http":
					url := args[0]
					dv["spec"].(map[string]interface{})["source"] = map[string]interface{}{"http": map[string]interface{}{"url": url}}

				case "s3":
					url := args[0]
					dv["spec"].(map[string]interface{})["source"] = map[string]interface{}{"s3": map[string]interface{}{"url": url}}
				case "pvc":
					namespace := args[0]
					name := args[1]
					dv["spec"].(map[string]interface{})["source"] = map[string]interface{}{
						"pvc": map[string]interface{}{
							"namespace": namespace,
							"name":      name}}
				}

				err = structToYamlFile(destinationFile, dv)
				Expect(err).ToNot(HaveOccurred())

				By(fmt.Sprint("Verifying kubectl create"))
				Eventually(func() bool {
					_, err := RunKubectlCommand(f, "create", "-f", destinationFile, "-n", f.Namespace.Name)
					if err != nil {
						return true
					}
					return false
				}, timeout, pollingInterval).Should(BeTrue())

			},
				table.Entry("[test_id:1321]fail with http source with invalid url format", "http", invalidURLFormat),
				table.Entry("[test_id:1322]fail with http source with empty url", "http", ""),
				table.Entry("[test_id:1323][crit:low]fail with s3 source with invalid url format", "s3", invalidURLFormat),
				table.Entry("[test_id:1324][crit:low]fail with s3 source with empty url", "s3", ""),
				table.Entry("[test_id:1325]fail with empty PVC source namespace", "pvc", "", "test-pvc"),
				table.Entry("[test_id:1326]fail with empty PVC source name", "pvc", "test", ""),
				table.Entry("fail with source PVC doesn't exist", "pvc", "test", "test-pvc"),
			)

			table.DescribeTable("with Datavolume PVC size should", func(size string) {

				By("Reading yaml file from: " + datavolumeTestFile)
				err := yamlFiletoStruct(datavolumeTestFile, &dv)
				Expect(err).ToNot(HaveOccurred())

				dv["spec"].(map[string]interface{})["pvc"].(map[string]interface{})["resources"].(map[string]interface{})["requests"].(map[string]interface{})["storage"] = size
				err = structToYamlFile(destinationFile, dv)
				Expect(err).ToNot(HaveOccurred())

				By(fmt.Sprint("Verifying kubectl apply"))
				Eventually(func() bool {
					_, err := RunKubectlCommand(f, "create", "-f", destinationFile, "-n", f.Namespace.Name)
					if err != nil {
						return true
					}
					return false
				}, timeout, pollingInterval).Should(BeTrue())
			},
				table.Entry("[test_id:1033]fail with zero PVC size", "0"),
				table.Entry("[test_id:1327]fail with negative PVC size", "-500m"),
				table.Entry("[test_id:1328]fail with invalid PVC size", "invalid_size"),
			)

		})
	})

	Context("DataVolume Already Exists", func() {
		BeforeEach(func() {
			dataVolume := utils.NewDataVolumeWithHTTPImport(dataVolumeName, "500Mi", validURL)

			dataVolume, err := utils.CreateDataVolumeFromDefinition(f.CdiClient, f.Namespace.Name, dataVolume)
			Expect(err).ToNot(HaveOccurred())
		})

		AfterEach(func() {
			err := utils.DeleteDataVolume(f.CdiClient, f.Namespace.Name, dataVolumeName)
			Expect(err).ToNot(HaveOccurred())
		})
		It("[test_id:1030]should fail creating an already existing DataVolume", func() {
			By(fmt.Sprint("Verifying kubectl create"))
			Eventually(func() bool {

				_, err := RunKubectlCommand(f, "create", "-f", datavolumeTestFile, "-n", f.Namespace.Name)
				if err != nil {
					return true
				}
				return false
			}, timeout, pollingInterval).Should(BeTrue())

		})
	})

	Context("DataVolume destination PVC", func() {
		BeforeEach(func() {
			pvc := utils.NewPVCDefinition(dataVolumeName, "50Mi", nil, nil)

			pvc, err := utils.CreatePVCFromDefinition(f.K8sClient, f.Namespace.Name, pvc)
			Expect(err).ToNot(HaveOccurred())
		})

		AfterEach(func() {
			pvc, err := f.K8sClient.CoreV1().PersistentVolumeClaims(f.Namespace.Name).Get(dataVolumeName, metav1.GetOptions{})
			Expect(err).ToNot(HaveOccurred())
			err = utils.DeletePVC(f.K8sClient, f.Namespace.Name, pvc)
			Expect(err).ToNot(HaveOccurred())
		})
		It("should fail creating a DataVolume with already existing destination pvc", func() {
			By(fmt.Sprint("Verifying kubectl create"))
			Eventually(func() bool {

				_, err := RunKubectlCommand(f, "create", "-f", datavolumeTestFile, "-n", f.Namespace.Name)
				if err != nil {
					return true
				}
				return false
			}, timeout, pollingInterval).Should(BeTrue())

		})
	})

	Context("when creating data volumes from manual manifests", func() {
		table.DescribeTable("with manifests Datavolume should", func(destinationFile string, expectError bool) {
			By(fmt.Sprint("Verifying kubectl apply"))
			_, err := RunKubectlCommand(f, "create", "-f", destinationFile, "-n", f.Namespace.Name)
			if expectError {
				Expect(err).To(HaveOccurred())
			} else {
				Expect(err).ToNot(HaveOccurred())
			}
		},
			table.Entry("fail with blank image source and contentType archive", "manifests/dvBlankArchive.yaml", true),
			table.Entry("fail with invalid contentType", "manifests/dvInvalidContentType.yaml", true),
			table.Entry("fail with missing source", "manifests/dvMissingSource.yaml", true),
			table.Entry("fail with multiple sources", "manifests/dvMultiSource.yaml", true),
			table.Entry("fail with invalid URL for http source", "manifests/dvInvalidURL.yaml", true),
			table.Entry("fail with invalid source PVC", "manifests/dvInvalidSourcePVC.yaml", true),
			table.Entry("succeed with valid source http", "manifests/datavolume.yaml", false),
			table.Entry("fail with missing PVC spec", "manifests/dvMissingPVCSpec.yaml", true),
			table.Entry("fail with missing resources spec", "manifests/dvMissingResourcesSpec.yaml", true),
			table.Entry("fail with 0 size PVC", "manifests/dv0SizePVC.yaml", true),
		)

	})
})

func yamlFiletoStruct(fileName string, o *map[string]interface{}) error {
	yamlFile, err := ioutil.ReadFile(fileName)
	if err != nil {
		return err
	}
	err = yaml.Unmarshal(yamlFile, o)
	if err != nil {
		return err
	}
	return nil

}

func structToYamlFile(fileName string, o interface{}) error {
	yamlOutput, err := yaml.Marshal(o)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(fileName, yamlOutput, 0644)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>