<commit_before>package helpers

import (
	"fmt"
	"strings"

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

func SetupReadOnlyOrgAndSpace() (string, string) {
	homeDir := SetHomeDir()
	SetAPI()
	LoginCF()
	orgName := NewOrgName()
	spaceName1 := NewSpaceName()
	spaceName2 := NewSpaceName()
	Eventually(CF("create-org", orgName)).Should(Exit(0))
	Eventually(CF("create-space", spaceName1, "-o", orgName)).Should(Exit(0))
	Eventually(CF("create-space", spaceName2, "-o", orgName)).Should(Exit(0))
	DestroyHomeDir(homeDir)
	return orgName, spaceName1
}

func CreateOrgAndSpace(org string, space string) {
	CreateOrg(org)
	TargetOrg(org)
	CreateSpace(space)
}

func CreateOrg(org string) {
	Eventually(CF("create-org", org)).Should(Exit(0))
}

func CreateSpace(space string) {
	Eventually(CF("create-space", space)).Should(Exit(0))
}

func GetOrgGUID(orgName string) string {
	session := CF("org", "--guid", orgName)
	Eventually(session).Should(Exit(0))
	return strings.TrimSpace(string(session.Out.Contents()))
}

func GetSpaceGUID(spaceName string) string {
	session := CF("space", "--guid", spaceName)
	Eventually(session).Should(Exit(0))
	return strings.TrimSpace(string(session.Out.Contents()))
}

func QuickDeleteOrg(orgName string) {
	guid := GetOrgGUID(orgName)
	url := fmt.Sprintf("/v2/organizations/%s?recursive=true&async=true", guid)
	session := CF("curl", "-X", "DELETE", url)
	Eventually(session).Should(Exit(0))
}

func QuickDeleteOrgIfExists(orgName string) {
	session := CF("org", "--guid", orgName)
	Eventually(session).Should(Exit())
	if session.ExitCode() != 0 {
		return
	}
	guid := strings.TrimSpace(string(session.Out.Contents()))
	url := fmt.Sprintf("/v2/organizations/%s?recursive=true&async=true", guid)
	session = CF("curl", "-X", "DELETE", url)
	Eventually(session).Should(Exit())
}

func QuickDeleteSpace(spaceName string) {
	guid := GetSpaceGUID(spaceName)
	url := fmt.Sprintf("/v2/spaces/%s?recursive=true&async=true", guid)
	session := CF("curl", "-X", "DELETE", url)
	Eventually(session).Should(Exit(0))
}
<commit_msg>add logic to ignore empty orgs<commit_after>package helpers

import (
	"fmt"
	"strings"

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

func SetupReadOnlyOrgAndSpace() (string, string) {
	homeDir := SetHomeDir()
	SetAPI()
	LoginCF()
	orgName := NewOrgName()
	spaceName1 := NewSpaceName()
	spaceName2 := NewSpaceName()
	Eventually(CF("create-org", orgName)).Should(Exit(0))
	Eventually(CF("create-space", spaceName1, "-o", orgName)).Should(Exit(0))
	Eventually(CF("create-space", spaceName2, "-o", orgName)).Should(Exit(0))
	DestroyHomeDir(homeDir)
	return orgName, spaceName1
}

func CreateOrgAndSpace(org string, space string) {
	CreateOrg(org)
	TargetOrg(org)
	CreateSpace(space)
}

func CreateOrg(org string) {
	Eventually(CF("create-org", org)).Should(Exit(0))
}

func CreateSpace(space string) {
	Eventually(CF("create-space", space)).Should(Exit(0))
}

func GetOrgGUID(orgName string) string {
	session := CF("org", "--guid", orgName)
	Eventually(session).Should(Exit(0))
	return strings.TrimSpace(string(session.Out.Contents()))
}

func GetSpaceGUID(spaceName string) string {
	session := CF("space", "--guid", spaceName)
	Eventually(session).Should(Exit(0))
	return strings.TrimSpace(string(session.Out.Contents()))
}

func QuickDeleteOrg(orgName string) {
	// If orgName is empty, the BeforeSuite has failed and attempting to delete
	// will produce a meaningless error.
	if orgName == "" {
		fmt.Println("Empty org name. Skipping deletion.")
		return
	}

	guid := GetOrgGUID(orgName)
	url := fmt.Sprintf("/v2/organizations/%s?recursive=true&async=true", guid)
	session := CF("curl", "-X", "DELETE", url)
	Eventually(session).Should(Exit(0))
}

func QuickDeleteOrgIfExists(orgName string) {
	session := CF("org", "--guid", orgName)
	Eventually(session).Should(Exit())
	if session.ExitCode() != 0 {
		return
	}
	guid := strings.TrimSpace(string(session.Out.Contents()))
	url := fmt.Sprintf("/v2/organizations/%s?recursive=true&async=true", guid)
	session = CF("curl", "-X", "DELETE", url)
	Eventually(session).Should(Exit())
}

func QuickDeleteSpace(spaceName string) {
	guid := GetSpaceGUID(spaceName)
	url := fmt.Sprintf("/v2/spaces/%s?recursive=true&async=true", guid)
	session := CF("curl", "-X", "DELETE", url)
	Eventually(session).Should(Exit(0))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"time"

	"github.com/ljfranklin/terraform-resource/models"
	"github.com/ljfranklin/terraform-resource/storage"
	"github.com/ljfranklin/terraform-resource/terraform"
)

func main() {

	if len(os.Args) < 2 {
		log.Fatalf("Expected path to sources as first arg")
	}
	sourceDir := os.Args[1]
	if err := os.Chdir(sourceDir); err != nil {
		log.Fatalf("Failed to access source dir '%s': %s", sourceDir, err)
	}
	tmpDir, err := ioutil.TempDir(os.TempDir(), "terraform-resource-out")
	if err != nil {
		log.Fatalf("Failed to create tmp dir at '%s'", os.TempDir())
	}
	defer os.RemoveAll(tmpDir)

	req := models.OutRequest{}
	if err = json.NewDecoder(os.Stdin).Decode(&req); err != nil {
		log.Fatalf("Failed to read OutRequest: %s", err)
	}

	storageKey := req.Source.Key
	if storageKey == "" {
		log.Fatalf("Must specify 'key' under resource.source")
	}

	driverType := req.Source.StorageDriver
	if driverType == "" {
		driverType = models.S3Driver
	}

	var storageDriver storage.Storage
	switch driverType {
	case models.S3Driver:
		storageDriver = storage.NewS3(
			req.Source.AccessKeyID,
			req.Source.SecretAccessKey,
			req.Source.RegionName,
			req.Source.Bucket,
		)
	default:
		supportedDrivers := []string{models.S3Driver}
		log.Fatalf("Unknown storage_driver '%s'. Supported drivers are: %v", driverType, strings.Join(supportedDrivers, ", "))
	}

	if req.Params.TerraformSource == "" {
		log.Fatalf("Must specify 'terraform_source' under put params")
	}

	stateFilePath := path.Join(tmpDir, "terraform.tfstate")
	client := terraform.Client{
		Source:        req.Params.TerraformSource,
		StateFilePath: stateFilePath,
	}

	version := ""
	metadata := []models.MetadataField{}

	if req.Params.Action == models.DestroyAction {
		stateFile, createErr := os.Create(stateFilePath)
		if createErr != nil {
			log.Fatalf("Failed to create state file at '%s': %s", stateFilePath, createErr)
		}
		defer stateFile.Close()

		err = storageDriver.Download(storageKey, stateFile)
		if err != nil {
			log.Fatalf("Failed to download state file: %s", err)
		}
		stateFile.Close()

		if err = client.Destroy(req.Params.TerraformVars); err != nil {
			log.Fatalf("Failed to run terraform destroy.\nError: %s", err)
		}

		err = storageDriver.Delete(storageKey)
		if err != nil {
			log.Fatalf("Failed to delete state file: %s", err)
		}

		version = time.Now().UTC().Format(time.RFC3339)
	} else {

		version, err = storageDriver.Version(storageKey)
		if err != nil {
			log.Fatalf("Failed to check for existing state file from '%s': %s", storageKey, err)
		}
		if version != "" {
			stateFile, createErr := os.Create(stateFilePath)
			if createErr != nil {
				log.Fatalf("Failed to create state file at '%s': %s", stateFilePath, createErr)
			}
			defer stateFile.Close()

			err = storageDriver.Download(storageKey, stateFile)
			if err != nil {
				log.Fatalf("Failed to download state file: %s", err)
			}
			stateFile.Close()
		}

		if err = client.Apply(req.Params.TerraformVars); err != nil {
			log.Fatalf("Failed to run terraform apply.\nError: %s", err)
		}
		stateFile, err := os.Open(stateFilePath)
		if err != nil {
			log.Fatalf("Failed to open state file at '%s'", stateFilePath)
		}
		defer stateFile.Close()

		err = storageDriver.Upload(storageKey, stateFile)
		if err != nil {
			log.Fatalf("Failed to upload state file: %s", err)
		}

		version, err = storageDriver.Version(storageKey)
		if err != nil {
			log.Fatalf("Failed to retrieve version from '%s': %s", storageKey, err)
		}
		if version == "" {
			log.Fatalf("Couldn't find state file at: %s", storageKey)
		}

		output, err := client.Output()
		if err != nil {
			log.Fatalf("Failed to terraform output.\nError: %s", err)
		}
		for key, value := range output {
			metadata = append(metadata, models.MetadataField{
				Name:  key,
				Value: value,
			})
		}
	}

	resp := models.OutResponse{
		Version: models.Version{
			Version: version,
		},
		Metadata: metadata,
	}

	if err := json.NewEncoder(os.Stdout).Encode(resp); err != nil {
		log.Fatalf("Failed to write OutResponse: %s", err)
	}
}
<commit_msg>Extract large 'out' method into functions<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"time"

	"github.com/ljfranklin/terraform-resource/models"
	"github.com/ljfranklin/terraform-resource/storage"
	"github.com/ljfranklin/terraform-resource/terraform"
)

func main() {

	if len(os.Args) < 2 {
		log.Fatalf("Expected path to sources as first arg")
	}
	sourceDir := os.Args[1]
	if err := os.Chdir(sourceDir); err != nil {
		log.Fatalf("Failed to access source dir '%s': %s", sourceDir, err)
	}
	tmpDir, err := ioutil.TempDir(os.TempDir(), "terraform-resource-out")
	if err != nil {
		log.Fatalf("Failed to create tmp dir at '%s'", os.TempDir())
	}
	defer os.RemoveAll(tmpDir)

	req := models.OutRequest{}
	if err = json.NewDecoder(os.Stdin).Decode(&req); err != nil {
		log.Fatalf("Failed to read OutRequest: %s", err)
	}

	storageKey := req.Source.Key
	if storageKey == "" {
		log.Fatal("Must specify 'key' under resource.source")
	}
	if req.Params.TerraformSource == "" {
		log.Fatal("Must specify 'terraform_source' under put params")
	}

	storageDriver, err := buildStorageDriver(req)
	if err != nil {
		log.Fatal(err.Error())
	}

	stateFilePath := path.Join(tmpDir, "terraform.tfstate")
	client := terraform.Client{
		Source:        req.Params.TerraformSource,
		StateFilePath: stateFilePath,
	}

	resp := models.OutResponse{}
	if req.Params.Action == models.DestroyAction {
		resp, err = performDestroy(stateFilePath, req, client, storageDriver)
	} else {
		resp, err = performApply(stateFilePath, req, client, storageDriver)
	}
	if err != nil {
		log.Fatalf("Failed to run terraform with action '%s': %s", req.Params.Action, err)
	}

	if err := json.NewEncoder(os.Stdout).Encode(resp); err != nil {
		log.Fatalf("Failed to write OutResponse: %s", err)
	}
}

func buildStorageDriver(req models.OutRequest) (storage.Storage, error) {
	driverType := req.Source.StorageDriver
	if driverType == "" {
		driverType = models.S3Driver
	}

	var storageDriver storage.Storage
	switch driverType {
	case models.S3Driver:
		storageDriver = storage.NewS3(
			req.Source.AccessKeyID,
			req.Source.SecretAccessKey,
			req.Source.RegionName,
			req.Source.Bucket,
		)
	default:
		supportedDrivers := []string{models.S3Driver}
		return nil, fmt.Errorf("Unknown storage_driver '%s'. Supported drivers are: %v", driverType, strings.Join(supportedDrivers, ", "))
	}

	return storageDriver, nil
}

func performApply(stateFilePath string, req models.OutRequest, client terraform.Client, storageDriver storage.Storage) (models.OutResponse, error) {
	var nilResponse models.OutResponse

	storageKey := req.Source.Key
	version, err := storageDriver.Version(storageKey)
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to check for existing state file from '%s': %s", storageKey, err)
	}
	if version != "" {
		stateFile, createErr := os.Create(stateFilePath)
		if createErr != nil {
			return nilResponse, fmt.Errorf("Failed to create state file at '%s': %s", stateFilePath, createErr)
		}
		defer stateFile.Close()

		err = storageDriver.Download(storageKey, stateFile)
		if err != nil {
			return nilResponse, fmt.Errorf("Failed to download state file: %s", err)
		}
		stateFile.Close()
	}

	if err = client.Apply(req.Params.TerraformVars); err != nil {
		return nilResponse, fmt.Errorf("Failed to run terraform apply.\nError: %s", err)
	}
	stateFile, err := os.Open(stateFilePath)
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to open state file at '%s'", stateFilePath)
	}
	defer stateFile.Close()

	err = storageDriver.Upload(storageKey, stateFile)
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to upload state file: %s", err)
	}

	version, err = storageDriver.Version(storageKey)
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to retrieve version from '%s': %s", storageKey, err)
	}
	if version == "" {
		return nilResponse, fmt.Errorf("Couldn't find state file at: %s", storageKey)
	}

	clientOutput, err := client.Output()
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to terraform output.\nError: %s", err)
	}

	metadata := []models.MetadataField{}
	for key, value := range clientOutput {
		metadata = append(metadata, models.MetadataField{
			Name:  key,
			Value: value,
		})
	}

	resp := models.OutResponse{
		Version: models.Version{
			Version: version,
		},
		Metadata: metadata,
	}
	return resp, nil
}

func performDestroy(stateFilePath string, req models.OutRequest, client terraform.Client, storageDriver storage.Storage) (models.OutResponse, error) {
	var nilResponse models.OutResponse
	stateFile, createErr := os.Create(stateFilePath)
	if createErr != nil {
		return nilResponse, fmt.Errorf("Failed to create state file at '%s': %s", stateFilePath, createErr)
	}
	defer stateFile.Close()

	err := storageDriver.Download(req.Source.Key, stateFile)
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to download state file: %s", err)
	}
	stateFile.Close()

	if err = client.Destroy(req.Params.TerraformVars); err != nil {
		return nilResponse, fmt.Errorf("Failed to run terraform destroy.\nError: %s", err)
	}

	err = storageDriver.Delete(req.Source.Key)
	if err != nil {
		return nilResponse, fmt.Errorf("Failed to delete state file: %s", err)
	}

	version := time.Now().UTC().Format(time.RFC3339)

	resp := models.OutResponse{
		Version: models.Version{
			Version: version,
		},
		Metadata: []models.MetadataField{},
	}
	return resp, nil
}
<|endoftext|>


<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"
	"io"
	"time"

	"github.com/uber/tchannel/golang/typed"
	"golang.org/x/net/context"
)

// beginCall begins an outbound call on the connection
func (c *Connection) beginCall(ctx context.Context, serviceName string, callOptions *CallOptions) (*OutboundCall, error) {
	if err := c.withStateRLock(func() error {
		switch c.state {
		case connectionActive, connectionStartClose, connectionInboundClosed:
			return nil
		case connectionClosed:
			return ErrConnectionClosed
		case connectionWaitingToRecvInitReq, connectionWaitingToSendInitReq, connectionWaitingToRecvInitRes:
			return ErrConnectionNotReady
		}

		return nil
	}); err != nil {
		return nil, err
	}

	deadline, _ := ctx.Deadline()
	timeToLive := deadline.Sub(time.Now())
	if timeToLive < 0 {
		return nil, ErrTimeout
	}

	requestID := c.NextMessageID()
	mex, err := c.outbound.newExchange(ctx, messageTypeCallReq, requestID, 512)
	if err != nil {
		return nil, err
	}

	headers := callHeaders{
		CallerName: c.localPeerInfo.ServiceName,
	}
	callOptions.setHeaders(headers)

	call := new(OutboundCall)
	call.mex = mex
	call.conn = c
	call.callReq = callReq{
		id:         requestID,
		Headers:    headers,
		Service:    serviceName,
		TimeToLive: timeToLive,
	}
	call.log = PrefixedLogger(fmt.Sprintf("Out%v-Call ", requestID), c.log)

	// TODO(mmihic): It'd be nice to do this without an fptr
	call.messageForFragment = func(initial bool) message {
		if initial {
			return &call.callReq
		}

		return new(callReqContinue)
	}

	call.contents = newFragmentingWriter(call, c.checksumType.New())
	span := CurrentSpan(ctx)
	if span != nil {
		call.callReq.Tracing = *span.NewChildSpan()
	} else {
		// TODO(mmihic): Potentially reject calls that are made outside a root context?
		call.callReq.Tracing.EnableTracing(false)
	}

	response := new(OutboundCallResponse)
	response.mex = mex
	response.log = PrefixedLogger(fmt.Sprintf("Out%v-Response ", requestID), c.log)
	response.messageForFragment = func(initial bool) message {
		if initial {
			return &response.callRes
		}

		return new(callResContinue)
	}
	response.contents = newFragmentingReader(response)
	call.response = response
	return call, nil
}

// handleCallRes handles an incoming call req message, forwarding the
// frame to the response channel waiting for it
func (c *Connection) handleCallRes(frame *Frame) {
	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}

// handleCallResContinue handles an incoming call res continue message,
// forwarding the frame to the response channel waiting for it
func (c *Connection) handleCallResContinue(frame *Frame) {
	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}

// An OutboundCall is an active call to a remote peer.  A client makes a call
// by calling BeginCall on the Channel, writing argument content via
// ArgWriter2() ArgWriter3(), and then reading reading response data via the
// ArgReader2() and ArgReader3() methods on the Response() object.
type OutboundCall struct {
	reqResWriter

	callReq  callReq
	response *OutboundCallResponse
}

// Response provides access to the call's response object, which can be used to
// read response arguments
func (call *OutboundCall) Response() *OutboundCallResponse {
	return call.response
}

// writeOperation writes the operation (arg1) to the call
func (call *OutboundCall) writeOperation(operation []byte) error {
	return NewArgWriter(call.arg1Writer()).Write(operation)
}

// Arg2Writer returns a WriteCloser that can be used to write the second argument.
// The returned writer must be closed once the write is complete.
func (call *OutboundCall) Arg2Writer() (io.WriteCloser, error) {
	return call.arg2Writer()
}

// Arg3Writer returns a WriteCloser that can be used to write the last argument.
// The returned writer must be closed once the write is complete.
func (call *OutboundCall) Arg3Writer() (io.WriteCloser, error) {
	return call.arg3Writer()
}

// An OutboundCallResponse is the response to an outbound call
type OutboundCallResponse struct {
	reqResReader

	callRes callRes
}

// ApplicationError returns true if the call resulted in an application level error
// TODO(mmihic): In current implementation, you must have called Arg2Reader before this
// method returns the proper value.  We should instead have this block until the first
// fragment is available, if the first fragment hasn't been received.
func (response *OutboundCallResponse) ApplicationError() bool {
	// TODO(mmihic): Wait for first fragment
	return response.callRes.ResponseCode == responseApplicationError
}

// Arg2Reader returns an io.ReadCloser to read the second argument.
// The ReadCloser must be closed once the argument has been read.
func (response *OutboundCallResponse) Arg2Reader() (io.ReadCloser, error) {
	var operation []byte
	if err := NewArgReader(response.arg1Reader()).Read(&operation); err != nil {
		return nil, err
	}

	return response.arg2Reader()
}

// Arg3Reader returns an io.ReadCloser to read the last argument.
// The ReadCloser must be closed once the argument has been read.
func (response *OutboundCallResponse) Arg3Reader() (io.ReadCloser, error) {
	return response.arg3Reader()
}

// handleError andles an error coming back from the peer. If the error is a
// protocol level error, the entire connection will be closed.  If the error is
// a request specific error, it will be written to the request's response
// channel and converted into a SystemError returned from the next reader or
// access call.
func (c *Connection) handleError(frame *Frame) {
	var errorMessage errorMessage
	rbuf := typed.NewReadBuffer(frame.SizedPayload())
	if err := errorMessage.read(rbuf); err != nil {
		c.log.Warnf("Unable to read Error frame from %s: %v", c.remotePeerInfo, err)
		c.connectionError(err)
		return
	}

	if errorMessage.errCode == ErrCodeProtocol {
		c.log.Warnf("Peer %s reported protocol error: %s", c.remotePeerInfo, errorMessage.message)
		c.connectionError(errorMessage.AsSystemError())
		return
	}

	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}
<commit_msg>Add Format to OutboundCallResponse<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 (
	"fmt"
	"io"
	"time"

	"github.com/uber/tchannel/golang/typed"
	"golang.org/x/net/context"
)

// beginCall begins an outbound call on the connection
func (c *Connection) beginCall(ctx context.Context, serviceName string, callOptions *CallOptions) (*OutboundCall, error) {
	if err := c.withStateRLock(func() error {
		switch c.state {
		case connectionActive, connectionStartClose, connectionInboundClosed:
			return nil
		case connectionClosed:
			return ErrConnectionClosed
		case connectionWaitingToRecvInitReq, connectionWaitingToSendInitReq, connectionWaitingToRecvInitRes:
			return ErrConnectionNotReady
		}

		return nil
	}); err != nil {
		return nil, err
	}

	deadline, _ := ctx.Deadline()
	timeToLive := deadline.Sub(time.Now())
	if timeToLive < 0 {
		return nil, ErrTimeout
	}

	requestID := c.NextMessageID()
	mex, err := c.outbound.newExchange(ctx, messageTypeCallReq, requestID, 512)
	if err != nil {
		return nil, err
	}

	headers := callHeaders{
		CallerName: c.localPeerInfo.ServiceName,
	}
	callOptions.setHeaders(headers)

	call := new(OutboundCall)
	call.mex = mex
	call.conn = c
	call.callReq = callReq{
		id:         requestID,
		Headers:    headers,
		Service:    serviceName,
		TimeToLive: timeToLive,
	}
	call.log = PrefixedLogger(fmt.Sprintf("Out%v-Call ", requestID), c.log)

	// TODO(mmihic): It'd be nice to do this without an fptr
	call.messageForFragment = func(initial bool) message {
		if initial {
			return &call.callReq
		}

		return new(callReqContinue)
	}

	call.contents = newFragmentingWriter(call, c.checksumType.New())
	span := CurrentSpan(ctx)
	if span != nil {
		call.callReq.Tracing = *span.NewChildSpan()
	} else {
		// TODO(mmihic): Potentially reject calls that are made outside a root context?
		call.callReq.Tracing.EnableTracing(false)
	}

	response := new(OutboundCallResponse)
	response.mex = mex
	response.log = PrefixedLogger(fmt.Sprintf("Out%v-Response ", requestID), c.log)
	response.messageForFragment = func(initial bool) message {
		if initial {
			return &response.callRes
		}

		return new(callResContinue)
	}
	response.contents = newFragmentingReader(response)
	call.response = response
	return call, nil
}

// handleCallRes handles an incoming call req message, forwarding the
// frame to the response channel waiting for it
func (c *Connection) handleCallRes(frame *Frame) {
	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}

// handleCallResContinue handles an incoming call res continue message,
// forwarding the frame to the response channel waiting for it
func (c *Connection) handleCallResContinue(frame *Frame) {
	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}

// An OutboundCall is an active call to a remote peer.  A client makes a call
// by calling BeginCall on the Channel, writing argument content via
// ArgWriter2() ArgWriter3(), and then reading reading response data via the
// ArgReader2() and ArgReader3() methods on the Response() object.
type OutboundCall struct {
	reqResWriter

	callReq  callReq
	response *OutboundCallResponse
}

// Response provides access to the call's response object, which can be used to
// read response arguments
func (call *OutboundCall) Response() *OutboundCallResponse {
	return call.response
}

// writeOperation writes the operation (arg1) to the call
func (call *OutboundCall) writeOperation(operation []byte) error {
	return NewArgWriter(call.arg1Writer()).Write(operation)
}

// Arg2Writer returns a WriteCloser that can be used to write the second argument.
// The returned writer must be closed once the write is complete.
func (call *OutboundCall) Arg2Writer() (io.WriteCloser, error) {
	return call.arg2Writer()
}

// Arg3Writer returns a WriteCloser that can be used to write the last argument.
// The returned writer must be closed once the write is complete.
func (call *OutboundCall) Arg3Writer() (io.WriteCloser, error) {
	return call.arg3Writer()
}

// An OutboundCallResponse is the response to an outbound call
type OutboundCallResponse struct {
	reqResReader

	callRes callRes
}

// ApplicationError returns true if the call resulted in an application level error
// TODO(mmihic): In current implementation, you must have called Arg2Reader before this
// method returns the proper value.  We should instead have this block until the first
// fragment is available, if the first fragment hasn't been received.
func (response *OutboundCallResponse) ApplicationError() bool {
	// TODO(mmihic): Wait for first fragment
	return response.callRes.ResponseCode == responseApplicationError
}

// Format the format of the request from the ArgScheme transport header.
func (response *OutboundCallResponse) Format() Format {
	return Format(response.callRes.Headers[ArgScheme])
}

// Arg2Reader returns an io.ReadCloser to read the second argument.
// The ReadCloser must be closed once the argument has been read.
func (response *OutboundCallResponse) Arg2Reader() (io.ReadCloser, error) {
	var operation []byte
	if err := NewArgReader(response.arg1Reader()).Read(&operation); err != nil {
		return nil, err
	}

	return response.arg2Reader()
}

// Arg3Reader returns an io.ReadCloser to read the last argument.
// The ReadCloser must be closed once the argument has been read.
func (response *OutboundCallResponse) Arg3Reader() (io.ReadCloser, error) {
	return response.arg3Reader()
}

// handleError andles an error coming back from the peer. If the error is a
// protocol level error, the entire connection will be closed.  If the error is
// a request specific error, it will be written to the request's response
// channel and converted into a SystemError returned from the next reader or
// access call.
func (c *Connection) handleError(frame *Frame) {
	var errorMessage errorMessage
	rbuf := typed.NewReadBuffer(frame.SizedPayload())
	if err := errorMessage.read(rbuf); err != nil {
		c.log.Warnf("Unable to read Error frame from %s: %v", c.remotePeerInfo, err)
		c.connectionError(err)
		return
	}

	if errorMessage.errCode == ErrCodeProtocol {
		c.log.Warnf("Peer %s reported protocol error: %s", c.remotePeerInfo, errorMessage.message)
		c.connectionError(errorMessage.AsSystemError())
		return
	}

	if err := c.outbound.forwardPeerFrame(frame); err != nil {
		c.outbound.removeExchange(frame.Header.ID)
	}
}
<|endoftext|>


<commit_before>package mat

import (
	"math"
	"math/rand"
)

func CopyOfMatrix(src [][]float64) (dst [][]float64) {
	dst = make([][]float64, len(src))
	for i, row := range src {
		dst[i] = make([]float64, len(row))
		copy(dst[i], row)
	}
	return
}

func CopyOfVector(src []float64) (dst []float64) {
	dst = make([]float64, len(src))
	copy(dst, src)
	return
}

func SumVector(dst, src []float64) {
	for i, val := range src {
		dst[i] += val
	}
}

func SumMatrix(dst, src [][]float64) {
	for i, row := range src {
		SumVector(dst[i], row)
	}
}

func SumVectors(dst []float64, srcs ...[]float64) {
	cols := len(dst)
	sum := 0.0
	for i := 0; i < cols; i++ {
		sum = 0
		for _, src := range srcs {
			sum += src[i]
		}
		dst[i] = sum
	}
}

func SumMatrixes(dst [][]float64, srcs ...[][]float64) {
	rowsCount := len(dst)
	rows := make([][]float64, len(srcs), len(srcs))
	for r := 0; r < rowsCount; r++ {
		for i, src := range srcs {
			rows[i] = src[r]
		}
		SumVectors(dst[r], rows...)
	}
}

func SubVector(dst, src []float64) {
	for i, val := range src {
		dst[i] -= val
	}
}

func SubMatrix(dst, src [][]float64) {
	for i, row := range src {
		SubVector(dst[i], row)
	}
}

func MulVectorByScalar(dst []float64, scalar float64) {
	for i, val := range dst {
		dst[i] = val * scalar
	}
}

func MulMatrixByScalar(dst [][]float64, scalar float64) {
	for _, row := range dst {
		MulVectorByScalar(row, scalar)
	}
}

func MulTransposeVector(dst [][]float64, a, b []float64) [][]float64 {
	if dst == nil {
		dst = make([][]float64, len(a))
		for i := range a {
			dst[i] = make([]float64, len(b))
		}
	}

	for i, valA := range a {
		row := dst[i]
		for j, valB := range b {
			row[j] = valA * valB
		}
	}
	return dst
}

func RandomVector(size int) []float64 {
	vector := make([]float64, size, size)
	for col := range vector {
		vector[col] = rand.NormFloat64()
	}

	return vector
}

func RandomMatrix(rows, cols int) [][]float64 {
	data := make([][]float64, rows)
	for row := range data {
		data[row] = RandomVector(cols)
	}
	return data
}

func MulVectorElementWise(dst, a, b []float64) []float64 {
	if dst == nil {
		dst = make([]float64, len(a), len(a))
	}
	for i := range dst {
		dst[i] = a[i] * b[i]
	}
	return dst
}

func SubVectorElementWise(a, b []float64) (diff []float64) {
	diff = make([]float64, len(a), len(a))
	for i := range diff {
		diff[i] = a[i] - b[i]
	}
	return
}

func VectorLen(a []float64) (vLen float64) {
	for _, val := range a {
		vLen += val * val
	}
	vLen = math.Sqrt(vLen)
	return
}

func ArgMax(a []float64) int {
	maxVal := math.SmallestNonzeroFloat64
	maxArg := -1

	for i, val := range a {
		if val > maxVal {
			maxVal = val
			maxArg = i
		}
	}
	return maxArg
}

func ZeroVector(a []float64) {
	for i := range a {
		a[i] = 0
	}
}

func ZeroMatrix(a [][]float64) {
	for _, vec := range a {
		ZeroVector(vec)
	}
}

func ZeroVectorOfMatrixes(a [][][]float64) {
	for _, m := range a {
		ZeroMatrix(m)
	}
}
<commit_msg>Mat doc strings<commit_after>package mat

import (
	"math"
	"math/rand"
)

// CopyOfMatrix makes a deep copy of matrix (slice of slices of float64)
func CopyOfMatrix(src [][]float64) (dst [][]float64) {
	dst = make([][]float64, len(src))
	for i, row := range src {
		dst[i] = make([]float64, len(row))
		copy(dst[i], row)
	}
	return
}

// CopyOfVector makes a deep copy of vector (slice of float64)
func CopyOfVector(src []float64) (dst []float64) {
	dst = make([]float64, len(src))
	copy(dst, src)
	return
}

// SumVector adds values from src slice to dst slice
func SumVector(dst, src []float64) {
	for i, val := range src {
		dst[i] += val
	}
}

// SumMatrix adds values from src matrix to dst matrix
func SumMatrix(dst, src [][]float64) {
	for i, row := range src {
		SumVector(dst[i], row)
	}
}

// SumVectors set dst values as sum of src slices
func SumVectors(dst []float64, srcs ...[]float64) {
	cols := len(dst)
	sum := 0.0
	for i := 0; i < cols; i++ {
		sum = 0
		for _, src := range srcs {
			sum += src[i]
		}
		dst[i] = sum
	}
}

// SumMatrixes set dst values as sum of src matrixes
func SumMatrixes(dst [][]float64, srcs ...[][]float64) {
	rowsCount := len(dst)
	rows := make([][]float64, len(srcs), len(srcs))
	for r := 0; r < rowsCount; r++ {
		for i, src := range srcs {
			rows[i] = src[r]
		}
		SumVectors(dst[r], rows...)
	}
}

// SubVector subtracts src value from dst
func SubVector(dst, src []float64) {
	for i, val := range src {
		dst[i] -= val
	}
}

// SubMatrix subtracts src value from dst
func SubMatrix(dst, src [][]float64) {
	for i, row := range src {
		SubVector(dst[i], row)
	}
}

// MulTransposeVector multiplies every values in dst by a constant factor
func MulVectorByScalar(dst []float64, scalar float64) {
	for i, val := range dst {
		dst[i] = val * scalar
	}
}

// MulMatrixByScalar multiplies every values in dst by a constant factor
func MulMatrixByScalar(dst [][]float64, scalar float64) {
	for _, row := range dst {
		MulVectorByScalar(row, scalar)
	}
}

// MulTransposeVector multiplies two matrices a' and b and places them to dst
func MulTransposeVector(dst [][]float64, a, b []float64) [][]float64 {
	if dst == nil {
		dst = make([][]float64, len(a))
		for i := range a {
			dst[i] = make([]float64, len(b))
		}
	}

	for i, valA := range a {
		row := dst[i]
		for j, valB := range b {
			row[j] = valA * valB
		}
	}
	return dst
}

// RandomVector creates vector of given size.
// Values are distributes using normal distribution
func RandomVector(size int) []float64 {
	vector := make([]float64, size, size)
	for col := range vector {
		vector[col] = rand.NormFloat64()
	}

	return vector
}

// RandomMatrix creates matrix of given size.
// Values are distributes using normal distribution
func RandomMatrix(rows, cols int) [][]float64 {
	data := make([][]float64, rows)
	for row := range data {
		data[row] = RandomVector(cols)
	}
	return data
}

// MulVectorElementWise multiplies a by b value by value.
// Result is set to dst
func MulVectorElementWise(dst, a, b []float64) []float64 {
	if dst == nil {
		dst = make([]float64, len(a), len(a))
	}
	for i := range dst {
		dst[i] = a[i] * b[i]
	}
	return dst
}

// SubVectorElementWise subtracts b from a (a-b).
// Result is retuned as dst
func SubVectorElementWise(a, b []float64) (diff []float64) {
	diff = make([]float64, len(a), len(a))
	for i := range diff {
		diff[i] = a[i] - b[i]
	}
	return
}

// VectorLen calculates euclidean length of the vector
func VectorLen(a []float64) (vLen float64) {
	for _, val := range a {
		vLen += val * val
	}
	vLen = math.Sqrt(vLen)
	return
}

// ArgMax calculates argmax(a)
func ArgMax(a []float64) int {
	maxVal := math.SmallestNonzeroFloat64
	maxArg := -1

	for i, val := range a {
		if val > maxVal {
			maxVal = val
			maxArg = i
		}
	}
	return maxArg
}

// ZeroVector sets all values to 0
func ZeroVector(a []float64) {
	for i := range a {
		a[i] = 0
	}
}

// ZeroMatrix sets all values to 0
func ZeroMatrix(a [][]float64) {
	for _, vec := range a {
		ZeroVector(vec)
	}
}

// ZeroVectorOfMatrixes sets all values to 0
func ZeroVectorOfMatrixes(a [][][]float64) {
	for _, m := range a {
		ZeroMatrix(m)
	}
}
<|endoftext|>


<commit_before>package chewing

import (
    "errors"
)

type PhraseBKForest struct {
    tree map[int]*PhraseBKTreeNode
}

type PhraseBKTreeNode struct {
    children map[uint8]*PhraseBKTreeNode
    phraseArrayItem *PhraseArrayItem
}

func NewPhraseBKForest() (phraseBKForest *PhraseBKForest) {
    phraseBKForest = new(PhraseBKForest)
    phraseBKForest.tree = make(map[int]*PhraseBKTreeNode)
    return phraseBKForest
}

func NewPhraseBKTreeNode(phraseArrayItem *PhraseArrayItem) (phraseBKTreeNode *PhraseBKTreeNode) {
    phraseBKTreeNode = new(PhraseBKTreeNode)
    phraseBKTreeNode.children = make(map[uint8]*PhraseBKTreeNode)
    phraseBKTreeNode.phraseArrayItem = phraseArrayItem
    return phraseBKTreeNode
}

func (this *PhraseBKForest) insert(phraseArrayItem *PhraseArrayItem) (err error) {
    length := len(phraseArrayItem.phoneSeq)
    if this.tree[length] == nil {
        this.tree[length] = NewPhraseBKTreeNode(phraseArrayItem)
        return nil
    } else {
        return this.tree[length].insert(phraseArrayItem)
    }
}

func (this *PhraseBKTreeNode) insert(phraseArrayItem *PhraseArrayItem) (err error) {
    distance, err := calculateHammingDistance(this.phraseArrayItem.phoneSeq, phraseArrayItem.phoneSeq)
    if err != nil {
        return err
    }
    if distance == 0 {
        return errors.New("Duplicate phoneSeq insert")
    }

    if this.children[distance] == nil {
        this.children[distance] = NewPhraseBKTreeNode(phraseArrayItem)
        return nil
    } else {
        return this.children[distance].insert(phraseArrayItem)
    }
}
<commit_msg>Implement BK-forest query<commit_after>package chewing

import (
    "errors"
)

type PhraseBKForest struct {
    tree map[int]*PhraseBKTreeNode
}

type PhraseBKTreeNode struct {
    children map[uint8]*PhraseBKTreeNode
    phraseArrayItem *PhraseArrayItem
}

func NewPhraseBKForest() (phraseBKForest *PhraseBKForest) {
    phraseBKForest = new(PhraseBKForest)
    phraseBKForest.tree = make(map[int]*PhraseBKTreeNode)
    return phraseBKForest
}

func NewPhraseBKTreeNode(phraseArrayItem *PhraseArrayItem) (phraseBKTreeNode *PhraseBKTreeNode) {
    phraseBKTreeNode = new(PhraseBKTreeNode)
    phraseBKTreeNode.children = make(map[uint8]*PhraseBKTreeNode)
    phraseBKTreeNode.phraseArrayItem = phraseArrayItem
    return phraseBKTreeNode
}

func (this *PhraseBKForest) insert(phraseArrayItem *PhraseArrayItem) (err error) {
    length := len(phraseArrayItem.phoneSeq)
    if this.tree[length] == nil {
        this.tree[length] = NewPhraseBKTreeNode(phraseArrayItem)
        return nil
    } else {
        return this.tree[length].insert(phraseArrayItem)
    }
}

func (this *PhraseBKTreeNode) insert(phraseArrayItem *PhraseArrayItem) (err error) {
    distance, err := calculateHammingDistance(this.phraseArrayItem.phoneSeq, phraseArrayItem.phoneSeq)
    if err != nil {
        return err
    }
    if distance == 0 {
        return errors.New("Duplicate phoneSeq insert")
    }

    if this.children[distance] == nil {
        this.children[distance] = NewPhraseBKTreeNode(phraseArrayItem)
        return nil
    } else {
        return this.children[distance].insert(phraseArrayItem)
    }
}

func (this *PhraseBKForest) query(phoneSeq []uint16, threshold uint8) (phraseArrayItem []*PhraseArrayItem) {
    length := len(phoneSeq)
    if this.tree[length] == nil {
        return make([]*PhraseArrayItem, 0)
    }

    phraseArrayItem = make([]*PhraseArrayItem, 0, 2)
    result := make(chan *PhraseArrayItem)
    count := make(chan int)

    counter := 1
    go this.tree[length].query(phoneSeq, threshold, count, result)

    for counter > 0 {
        select {
        case res := <- result:
            if len(phraseArrayItem) == cap(phraseArrayItem) {
                origin := phraseArrayItem
                phraseArrayItem = make([]*PhraseArrayItem, 0, len(origin))
                copy(phraseArrayItem, origin)
            }
            phraseArrayItem = phraseArrayItem[:len(phraseArrayItem) + 1]
            phraseArrayItem[len(phraseArrayItem) - 1] = res
        case c := <- count:
            counter += c
        }
    }

    return phraseArrayItem
}

func (this *PhraseBKTreeNode) query(phoneSeq []uint16, threshold uint8, count chan<- int, result chan<- *PhraseArrayItem) {
    diff, err := calculateHammingDistance(phoneSeq, this.phraseArrayItem.phoneSeq)
    if err != nil {
        panic("calculateHammingDistance fails in PhraseBKTreeNode.query")
    }
    if diff <= threshold {
        result <- this.phraseArrayItem
    }

    for i := diff - threshold; i <= diff + threshold; i++ {
        if this.children[i] != nil {
            count <- 1
            go this.children[i].query(phoneSeq, threshold, count, result)
        }
    }

    count <- -1
}
<|endoftext|>


<commit_before>package go_koans

import "fmt"

func aboutControlFlow() {
	{
		a, b, c := 1, 2, 3
		assert(a == __int__) // multiple assignment
		assert(b == __int__) // can make
		assert(c == __int__) // life easier
	}

	var str string

	{
		if 3.14 == 3 {
			str = "what is love?"
		} else {
			str = "baby dont hurt me"
		}
		assert(str == __string__) // no more

		if length := len(str); length == 17 {
			str = "to be"
		} else {
			str = "or not"
		}
		assert(str == __string__) // that is the question
	}

	{
		hola1, hola2 := "ho", "la"

		switch "hello" {
		case "hello":
			str = "hi"
		case "world":
			str = "planet"
		case fmt.Sprintf("%s%s", hola1, hola2):
			str = "senor"
		}
		assert(str == __string__) // cases can be of any type, even arbitrary expressions

		switch {
		case false:
			str = "first"
		case true:
			str = "second"
		}
		assert(str == __string__) // in the absence of value, there is truth
	}

	{
		n := 0
		for i := 0; i < 5; i++ {
			n += i
		}
		assert(n == __int__) // for can have the structure with which we are all familiar
	}

	{
		n := 1
		for {
			n *= 2
			if n > 20 {
				break
			}
		}
		assert(n == __int__) // though omitting everything creates an infinite loop
	}
}
<commit_msg>Control Koan<commit_after>package go_koans

import "fmt"

func aboutControlFlow() {
	{
		a, b, c := 1, 2, 3
		assert(a == 1) // multiple assignment
		assert(b == 2) // can make
		assert(c == 3) // life easier
	}

	var str string

	{
		if 3.14 == 3 {
			str = "what is love?"
		} else {
			str = "baby dont hurt me"
		}
		assert(str == "baby dont hurt me") // no more

		if length := len(str); length == 17 {
			str = "to be"
		} else {
			str = "or not"
		}
		assert(str == "to be") // that is the question
	}

	{
		hola1, hola2 := "ho", "la"

		switch "hello" {
		case "hello":
			str = "hi"
		case "world":
			str = "planet"
		case fmt.Sprintf("%s%s", hola1, hola2):
			str = "senor"
		}
		assert(str == "hi") // cases can be of any type, even arbitrary expressions

		switch {
		case false:
			str = "first"
		case true:
			str = "second"
		}
		assert(str == "second") // in the absence of value, there is truth
	}

	{
		n := 0
		for i := 0; i < 5; i++ {
			n += i
		}
		assert(n == 10) // for can have the structure with which we are all familiar
	}

	{
		n := 1
		for {
			n *= 2
			if n > 20 {
				break
			}
		}
		assert(n == 32) // though omitting everything creates an infinite loop
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package systests

import (
	"testing"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/service"
)

func TestPassphraseChange(t *testing.T) {
	tc := setupTest(t, "pp")
	tc2 := cloneContext(tc)

	libkb.G.LocalDb = nil

	defer tc.Cleanup()

	stopCh := make(chan error)
	svc := service.NewService(tc.G, false)
	startCh := svc.GetStartChannel()
	go func() {
		err := svc.Run()
		if err != nil {
			t.Logf("Running the service produced an error: %v", err)
		}
		stopCh <- err
	}()
	<-startCh

	userInfo := randomUser("pp")

	sui := signupUI{
		info:         userInfo,
		Contextified: libkb.NewContextified(tc2.G),
	}
	tc2.G.SetUI(&sui)
	signup := client.NewCmdSignupRunner(tc2.G)
	signup.SetTest()

	if err := signup.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(userInfo.passphrase); err != nil {
		t.Fatal(err)
	}

	oldPassphrase := userInfo.passphrase
	newPassphrase := userInfo.passphrase + userInfo.passphrase
	sui.info.passphrase = newPassphrase
	change := client.NewCmdPassphraseChangeRunner(tc2.G)

	if err := change.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(newPassphrase); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(oldPassphrase); err == nil {
		t.Fatal("old passphrase passed verification after passphrase change")
	}

	if err := client.CtlServiceStop(tc2.G); err != nil {
		t.Fatal(err)
	}

	// If the server failed, it's also an error
	if err := <-stopCh; err != nil {
		t.Fatal(err)
	}
}

func TestPassphraseRecover(t *testing.T) {
	tc := setupTest(t, "pp")
	tc2 := cloneContext(tc)

	libkb.G.LocalDb = nil

	defer tc.Cleanup()

	stopCh := make(chan error)
	svc := service.NewService(tc.G, false)
	startCh := svc.GetStartChannel()
	go func() {
		err := svc.Run()
		if err != nil {
			t.Logf("Running the service produced an error: %v", err)
		}
		stopCh <- err
	}()
	<-startCh

	userInfo := randomUser("pp")

	sui := signupUI{
		info:         userInfo,
		Contextified: libkb.NewContextified(tc2.G),
	}
	tc2.G.SetUI(&sui)
	signup := client.NewCmdSignupRunner(tc2.G)
	signup.SetTest()

	if err := signup.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(userInfo.passphrase); err != nil {
		t.Fatal(err)
	}

	// logout before recovering passphrase
	logout := client.NewCmdLogoutRunner(tc2.G)
	if err := logout.Run(); err != nil {
		t.Fatal(err)
	}

	// the paper key displayed during signup is in userInfo now, and it will be used
	// during passphrase recovery
	tc.G.Log.Debug("signup paper key: %s", userInfo.displayedPaperKey)

	oldPassphrase := userInfo.passphrase
	newPassphrase := userInfo.passphrase + userInfo.passphrase
	sui.info.passphrase = newPassphrase
	recoverCmd := client.NewCmdPassphraseRecoverRunner(tc2.G)

	if err := recoverCmd.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(newPassphrase); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(oldPassphrase); err == nil {
		t.Fatal("old passphrase passed verification after passphrase change")
	}

	if err := client.CtlServiceStop(tc2.G); err != nil {
		t.Fatal(err)
	}

	// If the server failed, it's also an error
	if err := <-stopCh; err != nil {
		t.Fatal(err)
	}
}
<commit_msg>recover test doesn't log out<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package systests

import (
	"testing"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/service"
)

func TestPassphraseChange(t *testing.T) {
	tc := setupTest(t, "pp")
	tc2 := cloneContext(tc)

	libkb.G.LocalDb = nil

	defer tc.Cleanup()

	stopCh := make(chan error)
	svc := service.NewService(tc.G, false)
	startCh := svc.GetStartChannel()
	go func() {
		err := svc.Run()
		if err != nil {
			t.Logf("Running the service produced an error: %v", err)
		}
		stopCh <- err
	}()
	<-startCh

	userInfo := randomUser("pp")

	sui := signupUI{
		info:         userInfo,
		Contextified: libkb.NewContextified(tc2.G),
	}
	tc2.G.SetUI(&sui)
	signup := client.NewCmdSignupRunner(tc2.G)
	signup.SetTest()

	if err := signup.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(userInfo.passphrase); err != nil {
		t.Fatal(err)
	}

	oldPassphrase := userInfo.passphrase
	newPassphrase := userInfo.passphrase + userInfo.passphrase
	sui.info.passphrase = newPassphrase
	change := client.NewCmdPassphraseChangeRunner(tc2.G)

	if err := change.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(newPassphrase); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(oldPassphrase); err == nil {
		t.Fatal("old passphrase passed verification after passphrase change")
	}

	if err := client.CtlServiceStop(tc2.G); err != nil {
		t.Fatal(err)
	}

	// If the server failed, it's also an error
	if err := <-stopCh; err != nil {
		t.Fatal(err)
	}
}

func TestPassphraseRecover(t *testing.T) {
	tc := setupTest(t, "pp")
	tc2 := cloneContext(tc)

	libkb.G.LocalDb = nil

	defer tc.Cleanup()

	stopCh := make(chan error)
	svc := service.NewService(tc.G, false)
	startCh := svc.GetStartChannel()
	go func() {
		err := svc.Run()
		if err != nil {
			t.Logf("Running the service produced an error: %v", err)
		}
		stopCh <- err
	}()
	<-startCh

	userInfo := randomUser("pp")

	sui := signupUI{
		info:         userInfo,
		Contextified: libkb.NewContextified(tc2.G),
	}
	tc2.G.SetUI(&sui)
	signup := client.NewCmdSignupRunner(tc2.G)
	signup.SetTest()

	if err := signup.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(userInfo.passphrase); err != nil {
		t.Fatal(err)
	}

	// the paper key displayed during signup is in userInfo now, and it will be used
	// during passphrase recovery
	tc.G.Log.Debug("signup paper key: %s", userInfo.displayedPaperKey)

	oldPassphrase := userInfo.passphrase
	newPassphrase := userInfo.passphrase + userInfo.passphrase
	sui.info.passphrase = newPassphrase
	recoverCmd := client.NewCmdPassphraseRecoverRunner(tc2.G)

	if err := recoverCmd.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(newPassphrase); err != nil {
		t.Fatal(err)
	}

	if _, err := tc.G.LoginState().VerifyPlaintextPassphrase(oldPassphrase); err == nil {
		t.Fatal("old passphrase passed verification after passphrase change")
	}

	if err := client.CtlServiceStop(tc2.G); err != nil {
		t.Fatal(err)
	}

	// If the server failed, it's also an error
	if err := <-stopCh; err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"unicode"

	"github.com/pkg/errors"
)

const (
	exitCodesDat = "https://raw.githubusercontent.com/arangodb/arangodb/main/lib/Basics/exitcodes.dat"
)

type exitCode struct {
	name        string
	code        int
	reason      string
	description string
}

func fatal(args ...interface{}) {
	fmt.Print(args...)
	os.Exit(1)
}

func main() {
	root, err := os.Getwd()
	if err != nil {
		fatal(err)
		return
	}

	dat, err := downloadArangodExitCodesDat()
	if err != nil {
		fatal(err)
		return
	}

	exitCodes, err := parseArangoDExitCodes(string(dat))
	if err != nil {
		fatal(err)
		return
	}

	err = generateExitCodesGoSource(exitCodes, root, err)
	if err != nil {
		fatal(err)
		return
	}

	fmt.Printf("ArangoD exit codes consts generated. Total %d codes found\n", len(exitCodes))
}

func generateExitCodesGoSource(exitCodes []exitCode, root string, err error) error {
	header, err := getLicenseHeader(root)
	if err != nil {
		return err
	}

	buf := bytes.Buffer{}
	buf.WriteString(header)
	buf.WriteString(`
// Code generated automatically. DO NOT EDIT.

package definitions

const (
`)
	for _, code := range exitCodes {
		buf.WriteString(fmt.Sprintf("	// %s\n", code.name))
		buf.WriteString(fmt.Sprintf(`	%s = %d // %s
`, getConstName(code.name), code.code, code.description))
	}
	buf.WriteString(")\n\n")

	buf.WriteString("var arangoDExitReason = map[int]string{\n")
	for _, code := range exitCodes {
		buf.WriteString(fmt.Sprintf("	// %s\n", code.name))
		buf.WriteString(fmt.Sprintf("	%s: \"%s\",\n", getConstName(code.name), code.reason))
	}
	buf.WriteString("}\n")

	resultFilePath := fmt.Sprintf("%s/pkg/definitions/exitcodes_generated.go", root)
	err = ioutil.WriteFile(resultFilePath, buf.Bytes(), 0600)
	return err
}

func getConstName(n string) string {
	prev := '_'
	// convert SOME_CONST_NAME to SomeConstName:
	dropUnderscoreAndStringify := func(r rune) rune {
		if prev == '_' {
			prev = r
			return unicode.ToTitle(r)
		}
		prev = r
		if r == '_' {
			return -1
		}
		return r
	}
	return "ArangoD" + strings.Map(dropUnderscoreAndStringify, strings.ToLower(n))
}

func downloadArangodExitCodesDat() ([]byte, error) {
	resp, err := http.Get(exitCodesDat)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}

func parseArangoDExitCodes(dat string) ([]exitCode, error) {
	// omit comments
	lines := strings.Split(dat, "\n")
	dat = ""
	for _, line := range lines {
		parts := strings.Split(line, "#")
		if len(parts[0]) > 0 {
			dat += parts[0] + "\n"
		}
	}

	b := bytes.NewBufferString(dat)
	csvReader := csv.NewReader(b)
	records, err := csvReader.ReadAll()
	if err != nil {
		return nil, err
	}
	result := make([]exitCode, 0, len(records))
	for _, r := range records {
		if len(r) < 4 {
			return nil, fmt.Errorf("expected at least 4 fields, got: %+v", r)
		}
		code, err := strconv.Atoi(r[1])
		if err != nil {
			return nil, errors.Wrapf(err, "while converting %s", r[1])
		}
		result = append(result, exitCode{
			name:        r[0],
			code:        code,
			reason:      r[2],
			description: r[3],
		})
	}

	sort.Slice(result, func(i, j int) bool {
		return result[i].code < result[j].code
	})

	return result, nil
}

func getLicenseHeader(root string) (string, error) {
	headerBoilerplate := fmt.Sprintf("%s/LICENSE.BOILERPLATE", root)
	b, err := ioutil.ReadFile(headerBoilerplate)
	if err != nil {
		return "", err
	}
	result := ""
	for _, line := range strings.Split(string(b), "\n") {
		if len(line) > 0 {
			result += "// " + line + "\n"
		} else {
			result += "//\n"
		}
	}
	return result, nil
}
<commit_msg>Use 3.10 branch instead of main<commit_after>//
// DISCLAIMER
//
// Copyright 2017-2022 ArangoDB GmbH, Cologne, Germany
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Copyright holder is ArangoDB GmbH, Cologne, Germany
//

package main

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"unicode"

	"github.com/pkg/errors"
)

const (
	// to support non-released version we have to use a version branch, not the main:
	arangoDVersion       = "3.10"
	exitCodesDatTemplate = "https://raw.githubusercontent.com/arangodb/arangodb/%s/lib/Basics/exitcodes.dat"
)

type exitCode struct {
	name        string
	code        int
	reason      string
	description string
}

func fatal(args ...interface{}) {
	fmt.Print(args...)
	os.Exit(1)
}

func main() {
	root, err := os.Getwd()
	if err != nil {
		fatal(err)
		return
	}

	dat, err := downloadArangodExitCodesDat()
	if err != nil {
		fatal(err)
		return
	}

	exitCodes, err := parseArangoDExitCodes(string(dat))
	if err != nil {
		fatal(err)
		return
	}

	err = generateExitCodesGoSource(exitCodes, root, err)
	if err != nil {
		fatal(err)
		return
	}

	fmt.Printf("ArangoD exit codes consts generated. Total %d codes found\n", len(exitCodes))
}

func generateExitCodesGoSource(exitCodes []exitCode, root string, err error) error {
	header, err := getLicenseHeader(root)
	if err != nil {
		return err
	}

	buf := bytes.Buffer{}
	buf.WriteString(header)
	buf.WriteString(`
// Code generated automatically. DO NOT EDIT.

package definitions

const (
`)
	for _, code := range exitCodes {
		buf.WriteString(fmt.Sprintf("	// %s\n", code.name))
		buf.WriteString(fmt.Sprintf(`	%s = %d // %s
`, getConstName(code.name), code.code, code.description))
	}
	buf.WriteString(")\n\n")

	buf.WriteString("var arangoDExitReason = map[int]string{\n")
	for _, code := range exitCodes {
		buf.WriteString(fmt.Sprintf("	// %s\n", code.name))
		buf.WriteString(fmt.Sprintf("	%s: \"%s\",\n", getConstName(code.name), code.reason))
	}
	buf.WriteString("}\n")

	resultFilePath := fmt.Sprintf("%s/pkg/definitions/exitcodes_generated.go", root)
	err = ioutil.WriteFile(resultFilePath, buf.Bytes(), 0600)
	return err
}

func getConstName(n string) string {
	prev := '_'
	// convert SOME_CONST_NAME to SomeConstName:
	dropUnderscoreAndStringify := func(r rune) rune {
		if prev == '_' {
			prev = r
			return unicode.ToTitle(r)
		}
		prev = r
		if r == '_' {
			return -1
		}
		return r
	}
	return "ArangoD" + strings.Map(dropUnderscoreAndStringify, strings.ToLower(n))
}

func downloadArangodExitCodesDat() ([]byte, error) {
	url := fmt.Sprintf(exitCodesDatTemplate, arangoDVersion)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	return ioutil.ReadAll(resp.Body)
}

func parseArangoDExitCodes(dat string) ([]exitCode, error) {
	// omit comments
	lines := strings.Split(dat, "\n")
	dat = ""
	for _, line := range lines {
		parts := strings.Split(line, "#")
		if len(parts[0]) > 0 {
			dat += parts[0] + "\n"
		}
	}

	b := bytes.NewBufferString(dat)
	csvReader := csv.NewReader(b)
	records, err := csvReader.ReadAll()
	if err != nil {
		return nil, err
	}
	result := make([]exitCode, 0, len(records))
	for _, r := range records {
		if len(r) < 4 {
			return nil, fmt.Errorf("expected at least 4 fields, got: %+v", r)
		}
		code, err := strconv.Atoi(r[1])
		if err != nil {
			return nil, errors.Wrapf(err, "while converting %s", r[1])
		}
		result = append(result, exitCode{
			name:        r[0],
			code:        code,
			reason:      r[2],
			description: r[3],
		})
	}

	sort.Slice(result, func(i, j int) bool {
		return result[i].code < result[j].code
	})

	return result, nil
}

func getLicenseHeader(root string) (string, error) {
	headerBoilerplate := fmt.Sprintf("%s/LICENSE.BOILERPLATE", root)
	b, err := ioutil.ReadFile(headerBoilerplate)
	if err != nil {
		return "", err
	}
	result := ""
	for _, line := range strings.Split(string(b), "\n") {
		if len(line) > 0 {
			result += "// " + line + "\n"
		} else {
			result += "//\n"
		}
	}
	return result, nil
}
<|endoftext|>


<commit_before>//Package runefrom provides a helper for converting an arbitrary string to a single rune.
package runefrom

import (
	"errors"
	"fmt"
	"unicode/utf8"
)

//String converts a nonempty string to a rune,
//returning an error if s contains more than one rune or the rune is invalid.
func String(s string) (rune, error) { //BUG(jmf): should handle \0 \n \t
	if s == "" {
		return -1, errors.New("no unicode code point specified")
	}
	r, sz := utf8.DecodeRuneInString(s)
	if r == utf8.RuneError {
		return -1, errors.New("invalid unicode code point")
	}
	if s[sz:] != "" {
		return -1, fmt.Errorf("expected a single unicode code point, got %q", s)
	}
	return r, nil
}
<commit_msg>handle escape sequences<commit_after>//Package runefrom provides a helper for converting an arbitrary string to a single rune.
package runefrom

import (
	"errors"
	"fmt"
	"strconv"
	"unicode/utf8"
)

//String converts a nonempty string to a rune,
//returning an error if s contains more than one rune or the rune is invalid.
func String(s string) (rune, error) {
	if s == "" {
		return -1, errors.New("no unicode code point specified")
	}
	r, _, t, err := strconv.UnquoteChar(s, 0)
	if err != nil {
		return -1, err
	}
	if r == utf8.RuneError {
		return -1, errors.New("invalid unicode code point")
	}
	if len(t) > 0 {
		return -1, fmt.Errorf("expected a single unicode code point, got %q", s)
	}
	return r, nil
}
<|endoftext|>


<commit_before>package partners

import (
	"fmt"
	"log"
	"net/url"
	"time"

	"github.com/brnstz/bus/internal/conf"
	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/bus/internal/models"

	"github.com/brnstz/bus/internal/partners/nyct_subway"
	"github.com/brnstz/bus/internal/partners/transit_realtime"

	"github.com/golang/protobuf/proto"
)

var (
	esiURL = "http://datamine.mta.info/mta_esi.php"

	routeToFeed = map[string]string{
		"1":  "1",
		"2":  "1",
		"3":  "1",
		"4":  "1",
		"5":  "1",
		"6":  "1",
		"6X": "1",
		"S":  "1",
		"GS": "1",
		"L":  "2",
		"SI": "11",
	}
)

func init() {
	log.Println("hello", nyct_subway.NyctTripDescriptor_Direction_name[1])
}

type mtaNYCSubway struct{}

func (_ mtaNYCSubway) Live(route models.Route, stop models.Stop) (d models.Departures, v []models.Vehicle, err error) {
	//now := time.Now()

	feed, exists := routeToFeed[stop.RouteID]
	if !exists {
		return
	}

	q := url.Values{}
	q.Set("key", conf.API.DatamineAPIKey)
	q.Set("feed_id", feed)
	u := fmt.Sprint(esiURL, "?", q.Encode())

	b, err := etc.RedisCache(u)
	if err != nil {
		log.Println("can't get live subways", err)
		return
	}

	tr := &transit_realtime.FeedMessage{}
	err = proto.Unmarshal(b, tr)
	if err != nil {
		log.Println("can't unmarshal", err)
		return
	}

	for _, e := range tr.Entity {
		//log.Println("what is it?", e)
		var vehicle models.Vehicle

		tripUpdate := e.GetTripUpdate()
		log.Println("what is trip?", tripUpdate.GetTrip().String())

		tripID := tripUpdate.GetTrip().GetTripId()

		updates := tripUpdate.GetStopTimeUpdate()

		first := true
		for _, u := range updates {
			//log.Println("what is update?", u)

			stopID := u.GetStopId()
			departureTime := time.Unix(u.GetDeparture().GetTime(), 0)

			// The first update in an entity is the stop where the train will
			// next be (including trips that haven't started yet)
			if first {
				first = false

				vehicle, err = models.GetVehicle(route.AgencyID, route.ID, stop.ID)
				vehicle.Live = true
				if err != nil {
					log.Println("can't get vehicle", err)
					return
				}
				v = append(v, vehicle)
			}

			// If this is our stop, then get the departure time.
			if stopID == stop.ID {
				d = append(d,
					&models.Departure{
						Time:   departureTime,
						TripID: tripID,
						Live:   true,
					},
				)
			}
		}
	}

	return
}
<commit_msg>live nyct subway vehicles<commit_after>package partners

import (
	"fmt"
	"log"
	"net/url"
	"time"

	"github.com/brnstz/bus/internal/conf"
	"github.com/brnstz/bus/internal/etc"
	"github.com/brnstz/bus/internal/models"

	"github.com/brnstz/bus/internal/partners/nyct_subway"
	"github.com/brnstz/bus/internal/partners/transit_realtime"

	"github.com/golang/protobuf/proto"
)

var (
	esiURL = "http://datamine.mta.info/mta_esi.php"

	routeToFeed = map[string]string{
		"1":  "1",
		"2":  "1",
		"3":  "1",
		"4":  "1",
		"5":  "1",
		"6":  "1",
		"6X": "1",
		"S":  "1",
		"GS": "1",
		"L":  "2",
		"SI": "11",
	}
)

func init() {
	log.Println("hello", nyct_subway.NyctTripDescriptor_Direction_name[1])
}

type mtaNYCSubway struct{}

func (_ mtaNYCSubway) Live(route models.Route, stop models.Stop) (d models.Departures, v []models.Vehicle, err error) {
	//now := time.Now()

	feed, exists := routeToFeed[stop.RouteID]
	if !exists {
		return
	}

	q := url.Values{}
	q.Set("key", conf.API.DatamineAPIKey)
	q.Set("feed_id", feed)
	u := fmt.Sprint(esiURL, "?", q.Encode())

	b, err := etc.RedisCache(u)
	if err != nil {
		log.Println("can't get live subways", err)
		return
	}

	tr := &transit_realtime.FeedMessage{}
	err = proto.Unmarshal(b, tr)
	if err != nil {
		log.Println("can't unmarshal", err)
		return
	}

	for _, e := range tr.Entity {
		var event interface{}

		var vehicle models.Vehicle

		tripUpdate := e.GetTripUpdate()
		trip := tripUpdate.GetTrip()
		if trip == nil {
			log.Println("skipping nil trip", e)
			continue
		}

		event, err = proto.GetExtension(trip, nyct_subway.E_NyctTripDescriptor)
		if err != nil {
			log.Println("can't get extension", err)
			continue
		}

		nycTrip, ok := event.(*nyct_subway.NyctTripDescriptor)
		if !ok {
			log.Println("can't coerce to nyct_subway.NyctTripDescriptor")
			continue
		}

		updates := tripUpdate.GetStopTimeUpdate()

		first := true
		for _, u := range updates {

			stopID := u.GetStopId()
			departureTime := time.Unix(u.GetDeparture().GetTime(), 0)

			// The first update in an entity is the stop where the train will
			// next be. Include only "assigned" trips, which are those that
			// are about to start.
			if first && nycTrip.GetIsAssigned() {
				first = false

				vehicle, err = models.GetVehicle(route.AgencyID, route.ID, stop.ID)
				vehicle.Live = true
				if err != nil {
					log.Println("can't get vehicle", err)
					return
				}
				v = append(v, vehicle)
			} else {
				first = false
			}

			// If this is our stop, then get the departure time.
			if stopID == stop.ID {
				d = append(d,
					&models.Departure{
						Time:   departureTime,
						TripID: trip.GetTripId(),
						Live:   true,
					},
				)
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 avro

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	avro "github.com/elodina/go-avro"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	var id int32
	if id, exists = schemaIdMap[schema]; exists {
		return id, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", SCHEMA_REGISTRY_V1_JSON)
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<commit_msg>Minor fixes.<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 avro

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	avro "github.com/elodina/go-avro"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	var id int32
	if id, exists = schemaIdMap[schema]; exists {
		return id, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if err := this.handleSuccess(response, decodedResponse); err != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", SCHEMA_REGISTRY_V1_JSON)
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 go_kafka_client

import (
	"encoding/json"
	"fmt"
	avro "github.com/stealthly/go-avro"
	"io"
	"net/http"
	"strconv"
	"strings"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	var id int32
	if id, exists = schemaIdMap[schema]; exists {
		return id, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}
		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	println(url)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes := make([]byte, response.ContentLength)
	response.Body.Read(responseBytes)
	println(string(responseBytes))
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes := make([]byte, response.ContentLength)
	response.Body.Read(responseBytes)
	err := json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<commit_msg>re #17 added cache feature for schema registry client.<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 go_kafka_client

import (
	"encoding/json"
	"fmt"
	avro "github.com/stealthly/go-avro"
	"io"
	"net/http"
	"strconv"
	"strings"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	var id int32
	if id, exists = schemaIdMap[schema]; exists {
		return id, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	println(url)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes := make([]byte, response.ContentLength)
	response.Body.Read(responseBytes)
	println(string(responseBytes))
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes := make([]byte, response.ContentLength)
	response.Body.Read(responseBytes)
	err := json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<|endoftext|>