<commit_before>package aws_client

import (
	"errors"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/aws/awsutil"
	"github.com/awslabs/aws-sdk-go/service/ec2"

	"github.com/xingzhou/go_service_broker/utils"
)

const (
	AMI_ID                = "ami-dc5e75b4" //"ami-ecb68a84"
	SECURITY_GROUP_ID     = "sg-b23aead6"
	SUBNET_ID             = "subnet-0c75a427"
	KEYPAIR_NAME          = "broker_keypair"
	INSTANCE_TYPE         = "t2.micro"
	LINUX_USER            = "ubuntu"
	KEYPAIR_DIR_NAME      = ".gsb"
	PIRVATE_KEY_FILE_NAME = "broker_id_rsa"
	PUBLIC_KEY_FILE_NAME  = "broker_id_rsa.pub"
)

type Client interface {
	CreateInstance() (string, error)
	GetInstanceState(instanceId string) (string, error)
	InjectKeyPair(instanceId string) (string, error)
	DeleteInstance(instanceId string) error
	RevokeKeyPair(instanceId string, privateKey string) error
}

type AWSClient struct {
	EC2Client *ec2.EC2
}

func NewClient(region string) *AWSClient {
	return &AWSClient{
		EC2Client: ec2.New(&aws.Config{Region: region}),
	}
}

func (c *AWSClient) CreateInstance() (string, error) {
	return c.createInstance(AMI_ID)
}

func (c *AWSClient) GetInstanceState(instanceId string) (string, error) {
	instanceInput := &ec2.DescribeInstancesInput{
		InstanceIDs: []*string{
			aws.String(instanceId), // Required
		},
	}

	instanceOutput, err := c.EC2Client.DescribeInstances(instanceInput)
	if err != nil {
		return "", err
	}

	state, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Reservations[0].Instances[0].State.Name))
	return state, nil
}

func (c *AWSClient) setupKeyPair() error {
	private_key_file := path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME, PIRVATE_KEY_FILE_NAME)

	if !utils.Exists(private_key_file) {
		keypairInput := &ec2.CreateKeyPairInput{
			KeyName: aws.String(KEYPAIR_NAME),
		}

		keypairOutput, err := c.EC2Client.CreateKeyPair(keypairInput)
		if err != nil {
			return err
		}

		key_dir := path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME)
		if _, err := os.Stat(key_dir); os.IsNotExist(err) {
			err = os.Mkdir(key_dir, 0700)
			if err != nil {
				return errors.New("failed to create local keypair directory")
			}
		}

		key_data, _ := strconv.Unquote(awsutil.StringValue(keypairOutput.KeyMaterial))
		err = utils.WriteFile(private_key_file, []byte(key_data))
		if err != nil {
			return errors.New("failed to save private key file")
		}
	}

	return nil
}

func (c *AWSClient) InjectKeyPair(instanceId string) (string, error) {
	instanceInput := &ec2.DescribeInstancesInput{
		InstanceIDs: []*string{
			aws.String(instanceId), // Required
		},
	}

	instanceOutput, err := c.EC2Client.DescribeInstances(instanceInput)
	if err != nil {
		return "", err
	}

	ip, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Reservations[0].Instances[0].PublicIPAddress))
	pemBytes, err := utils.ReadFile(path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME, PIRVATE_KEY_FILE_NAME))
	if err != nil {
		return "", err
	}

	awsSShClient, err := utils.GetSshClient(LINUX_USER, pemBytes, ip)
	if err != nil {
		return "", err
	}

	command := `rm -f ./broker_id_rsa ./broker_id_rsa.pub
		ssh-keygen -q -t rsa -N ""  -f ./broker_id_rsa
		cat ./broker_id_rsa.pub >> .ssh/authorized_keys
		cat ./broker_id_rsa`

	privateKey, err := awsSShClient.ExecCommand(command)
	if err != nil {
		return "", err
	}

	return privateKey, nil
}

func (c *AWSClient) createInstance(imageId string) (string, error) {
	err := c.setupKeyPair()
	if err != nil {
		return "", err
	}

	instanceInput := &ec2.RunInstancesInput{
		ImageID:  aws.String(imageId), // Required
		MaxCount: aws.Long(1),         // Required
		MinCount: aws.Long(1),         // Required
		// AdditionalInfo: aws.String("String"),
		// BlockDeviceMappings: []*ec2.BlockDeviceMapping{
		// 	&ec2.BlockDeviceMapping{ // Required
		// 		DeviceName: aws.String("String"),
		// 		EBS: &ec2.EBSBlockDevice{
		// 			DeleteOnTermination: aws.Boolean(true),
		// 			Encrypted:           aws.Boolean(true),
		// 			IOPS:                aws.Long(1),
		// 			SnapshotID:          aws.String("String"),
		// 			VolumeSize:          aws.Long(1),
		// 			VolumeType:          aws.String("VolumeType"),
		// 		},
		// 		NoDevice:    aws.String("String"),
		// 		VirtualName: aws.String("String"),
		// 	},
		// 	// More values...
		// },
		// ClientToken: aws.String("String"),
		// DisableAPITermination: aws.Boolean(true),
		// DryRun:                aws.Boolean(true),
		// EBSOptimized:          aws.Boolean(true),
		// IAMInstanceProfile: &ec2.IAMInstanceProfileSpecification{
		// 	ARN:  aws.String("String"),
		// 	Name: aws.String("String"),
		// },
		// InstanceInitiatedShutdownBehavior: aws.String("ShutdownBehavior"),
		InstanceType: aws.String(INSTANCE_TYPE),
		// KernelID:                          aws.String("String"),
		KeyName: aws.String(KEYPAIR_NAME),
		// Monitoring: &ec2.RunInstancesMonitoringEnabled{
		// 	Enabled: aws.Boolean(true), // Required
		// },
		// NetworkInterfaces: []*ec2.InstanceNetworkInterfaceSpecification{
		// 	&ec2.InstanceNetworkInterfaceSpecification{ // Required
		// 		AssociatePublicIPAddress: aws.Boolean(true),
		// 		DeleteOnTermination:      aws.Boolean(true),
		// 		Description:              aws.String("String"),
		// 		DeviceIndex:              aws.Long(1),
		// 		Groups: []*string{
		// 			aws.String("String"), // Required
		// 			// More values...
		// 		},
		// 		NetworkInterfaceID: aws.String("String"),
		// 		PrivateIPAddress:   aws.String("String"),
		// 		PrivateIPAddresses: []*ec2.PrivateIPAddressSpecification{
		// 			&ec2.PrivateIPAddressSpecification{ // Required
		// 				PrivateIPAddress: aws.String("String"), // Required
		// 				Primary:          aws.Boolean(true),
		// 			},
		// 			// More values...
		// 		},
		// 		SecondaryPrivateIPAddressCount: aws.Long(1),
		// 		SubnetID:                       aws.String("String"),
		// 	},
		// 	// More values...
		// },
		// Placement: &ec2.Placement{
		// 	AvailabilityZone: aws.String("String"),
		// 	GroupName:        aws.String("String"),
		// 	Tenancy:          aws.String("Tenancy"),
		// },
		// PrivateIPAddress: aws.String("String"),
		// RAMDiskID:        aws.String("String"),
		SecurityGroupIDs: []*string{
			aws.String(SECURITY_GROUP_ID), // Required
			// More values...
		},
		SubnetID: aws.String(SUBNET_ID),
	}

	instanceOutput, err := c.EC2Client.RunInstances(instanceInput)
	if err != nil {
		return "", err
	}

	fmt.Println(awsutil.StringValue(instanceOutput))
	instanceId, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Instances[0].InstanceID))

	return instanceId, nil
}

func (c *AWSClient) DeleteInstance(instanceId string) error {
	terminateInstanceInput := &ec2.TerminateInstancesInput{
		// One or more instance IDs.
		InstanceIDs: []*string{
			aws.String(instanceId), // Required
		},
	}

	_, err := c.EC2Client.TerminateInstances(terminateInstanceInput)
	if err != nil {
		return err
	}

	return nil
}

func (c *AWSClient) RevokeKeyPair(instanceId string, privateKey string) error {
	instanceInput := &ec2.DescribeInstancesInput{
		InstanceIDs: []*string{
			aws.String(instanceId),
		},
	}

	instanceOutput, err := c.EC2Client.DescribeInstances(instanceInput)
	if err != nil {
		return err
	}

	ip, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Reservations[0].Instances[0].PublicIPAddress))
	pemBytes, err := utils.ReadFile(path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME, PIRVATE_KEY_FILE_NAME))
	if err != nil {
		return err
	}

	awsSShClient, err := utils.GetSshClient(LINUX_USER, pemBytes, ip)
	if err != nil {
		return err
	}

	publicKey, err := utils.GeneratePublicKey([]byte(privateKey))
	if err != nil {
		return err
	}

	escapedPublicKey := strings.Replace(publicKey, "/", "\\/", -1)
	command := fmt.Sprintf("sed '/%s/d' -i ~/.ssh/authorized_keys && echo 'revoked the public key: %s'", escapedPublicKey, publicKey)

	result, err := awsSShClient.ExecCommand(command)
	if err != nil {
		return err
	}
	fmt.Println(result)

	return nil
}
<commit_msg>refactor to use utils MkDir<commit_after>package aws_client

import (
	"errors"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/aws/awsutil"
	"github.com/awslabs/aws-sdk-go/service/ec2"

	"github.com/xingzhou/go_service_broker/utils"
)

const (
	AMI_ID                = "ami-dc5e75b4" //"ami-ecb68a84"
	SECURITY_GROUP_ID     = "sg-b23aead6"
	SUBNET_ID             = "subnet-0c75a427"
	KEYPAIR_NAME          = "broker_keypair"
	INSTANCE_TYPE         = "t2.micro"
	LINUX_USER            = "ubuntu"
	KEYPAIR_DIR_NAME      = ".gsb"
	PIRVATE_KEY_FILE_NAME = "broker_id_rsa"
	PUBLIC_KEY_FILE_NAME  = "broker_id_rsa.pub"
)

type Client interface {
	CreateInstance() (string, error)
	GetInstanceState(instanceId string) (string, error)
	InjectKeyPair(instanceId string) (string, error)
	DeleteInstance(instanceId string) error
	RevokeKeyPair(instanceId string, privateKey string) error
}

type AWSClient struct {
	EC2Client *ec2.EC2
}

func NewClient(region string) *AWSClient {
	return &AWSClient{
		EC2Client: ec2.New(&aws.Config{Region: region}),
	}
}

func (c *AWSClient) CreateInstance() (string, error) {
	return c.createInstance(AMI_ID)
}

func (c *AWSClient) GetInstanceState(instanceId string) (string, error) {
	instanceInput := &ec2.DescribeInstancesInput{
		InstanceIDs: []*string{
			aws.String(instanceId), // Required
		},
	}

	instanceOutput, err := c.EC2Client.DescribeInstances(instanceInput)
	if err != nil {
		return "", err
	}

	state, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Reservations[0].Instances[0].State.Name))
	return state, nil
}

func (c *AWSClient) setupKeyPair() error {
	private_key_file := path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME, PIRVATE_KEY_FILE_NAME)

	if !utils.Exists(private_key_file) {
		keypairInput := &ec2.CreateKeyPairInput{
			KeyName: aws.String(KEYPAIR_NAME),
		}

		keypairOutput, err := c.EC2Client.CreateKeyPair(keypairInput)
		if err != nil {
			return err
		}

		key_dir := path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME)
		if !utils.MkDir(key_dir) {
			return errors.New("failed to create local keypair directory")
		}

		key_data, _ := strconv.Unquote(awsutil.StringValue(keypairOutput.KeyMaterial))
		err = utils.WriteFile(private_key_file, []byte(key_data))
		if err != nil {
			return errors.New("failed to save private key file")
		}
	}

	return nil
}

func (c *AWSClient) InjectKeyPair(instanceId string) (string, error) {
	instanceInput := &ec2.DescribeInstancesInput{
		InstanceIDs: []*string{
			aws.String(instanceId), // Required
		},
	}

	instanceOutput, err := c.EC2Client.DescribeInstances(instanceInput)
	if err != nil {
		return "", err
	}

	ip, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Reservations[0].Instances[0].PublicIPAddress))
	pemBytes, err := utils.ReadFile(path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME, PIRVATE_KEY_FILE_NAME))
	if err != nil {
		return "", err
	}

	awsSShClient, err := utils.GetSshClient(LINUX_USER, pemBytes, ip)
	if err != nil {
		return "", err
	}

	command := `rm -f ./broker_id_rsa ./broker_id_rsa.pub
		ssh-keygen -q -t rsa -N ""  -f ./broker_id_rsa
		cat ./broker_id_rsa.pub >> .ssh/authorized_keys
		cat ./broker_id_rsa`

	privateKey, err := awsSShClient.ExecCommand(command)
	if err != nil {
		return "", err
	}

	return privateKey, nil
}

func (c *AWSClient) createInstance(imageId string) (string, error) {
	err := c.setupKeyPair()
	if err != nil {
		return "", err
	}

	instanceInput := &ec2.RunInstancesInput{
		ImageID:  aws.String(imageId), // Required
		MaxCount: aws.Long(1),         // Required
		MinCount: aws.Long(1),         // Required
		// AdditionalInfo: aws.String("String"),
		// BlockDeviceMappings: []*ec2.BlockDeviceMapping{
		// 	&ec2.BlockDeviceMapping{ // Required
		// 		DeviceName: aws.String("String"),
		// 		EBS: &ec2.EBSBlockDevice{
		// 			DeleteOnTermination: aws.Boolean(true),
		// 			Encrypted:           aws.Boolean(true),
		// 			IOPS:                aws.Long(1),
		// 			SnapshotID:          aws.String("String"),
		// 			VolumeSize:          aws.Long(1),
		// 			VolumeType:          aws.String("VolumeType"),
		// 		},
		// 		NoDevice:    aws.String("String"),
		// 		VirtualName: aws.String("String"),
		// 	},
		// 	// More values...
		// },
		// ClientToken: aws.String("String"),
		// DisableAPITermination: aws.Boolean(true),
		// DryRun:                aws.Boolean(true),
		// EBSOptimized:          aws.Boolean(true),
		// IAMInstanceProfile: &ec2.IAMInstanceProfileSpecification{
		// 	ARN:  aws.String("String"),
		// 	Name: aws.String("String"),
		// },
		// InstanceInitiatedShutdownBehavior: aws.String("ShutdownBehavior"),
		InstanceType: aws.String(INSTANCE_TYPE),
		// KernelID:                          aws.String("String"),
		KeyName: aws.String(KEYPAIR_NAME),
		// Monitoring: &ec2.RunInstancesMonitoringEnabled{
		// 	Enabled: aws.Boolean(true), // Required
		// },
		// NetworkInterfaces: []*ec2.InstanceNetworkInterfaceSpecification{
		// 	&ec2.InstanceNetworkInterfaceSpecification{ // Required
		// 		AssociatePublicIPAddress: aws.Boolean(true),
		// 		DeleteOnTermination:      aws.Boolean(true),
		// 		Description:              aws.String("String"),
		// 		DeviceIndex:              aws.Long(1),
		// 		Groups: []*string{
		// 			aws.String("String"), // Required
		// 			// More values...
		// 		},
		// 		NetworkInterfaceID: aws.String("String"),
		// 		PrivateIPAddress:   aws.String("String"),
		// 		PrivateIPAddresses: []*ec2.PrivateIPAddressSpecification{
		// 			&ec2.PrivateIPAddressSpecification{ // Required
		// 				PrivateIPAddress: aws.String("String"), // Required
		// 				Primary:          aws.Boolean(true),
		// 			},
		// 			// More values...
		// 		},
		// 		SecondaryPrivateIPAddressCount: aws.Long(1),
		// 		SubnetID:                       aws.String("String"),
		// 	},
		// 	// More values...
		// },
		// Placement: &ec2.Placement{
		// 	AvailabilityZone: aws.String("String"),
		// 	GroupName:        aws.String("String"),
		// 	Tenancy:          aws.String("Tenancy"),
		// },
		// PrivateIPAddress: aws.String("String"),
		// RAMDiskID:        aws.String("String"),
		SecurityGroupIDs: []*string{
			aws.String(SECURITY_GROUP_ID), // Required
			// More values...
		},
		SubnetID: aws.String(SUBNET_ID),
	}

	instanceOutput, err := c.EC2Client.RunInstances(instanceInput)
	if err != nil {
		return "", err
	}

	fmt.Println(awsutil.StringValue(instanceOutput))
	instanceId, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Instances[0].InstanceID))

	return instanceId, nil
}

func (c *AWSClient) DeleteInstance(instanceId string) error {
	terminateInstanceInput := &ec2.TerminateInstancesInput{
		// One or more instance IDs.
		InstanceIDs: []*string{
			aws.String(instanceId), // Required
		},
	}

	_, err := c.EC2Client.TerminateInstances(terminateInstanceInput)
	if err != nil {
		return err
	}

	return nil
}

func (c *AWSClient) RevokeKeyPair(instanceId string, privateKey string) error {
	instanceInput := &ec2.DescribeInstancesInput{
		InstanceIDs: []*string{
			aws.String(instanceId),
		},
	}

	instanceOutput, err := c.EC2Client.DescribeInstances(instanceInput)
	if err != nil {
		return err
	}

	ip, _ := strconv.Unquote(awsutil.StringValue(instanceOutput.Reservations[0].Instances[0].PublicIPAddress))
	pemBytes, err := utils.ReadFile(path.Join(os.Getenv("HOME"), KEYPAIR_DIR_NAME, PIRVATE_KEY_FILE_NAME))
	if err != nil {
		return err
	}

	awsSShClient, err := utils.GetSshClient(LINUX_USER, pemBytes, ip)
	if err != nil {
		return err
	}

	publicKey, err := utils.GeneratePublicKey([]byte(privateKey))
	if err != nil {
		return err
	}

	escapedPublicKey := strings.Replace(publicKey, "/", "\\/", -1)
	command := fmt.Sprintf("sed '/%s/d' -i ~/.ssh/authorized_keys && echo 'revoked the public key: %s'", escapedPublicKey, publicKey)

	result, err := awsSShClient.ExecCommand(command)
	if err != nil {
		return err
	}
	fmt.Println(result)

	return nil
}
<|endoftext|>


<commit_before>package awsiam

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"text/template"

	"code.cloudfoundry.org/lager"
	awsarn "github.com/aws/aws-sdk-go/aws/arn"
	iampolicy "github.com/minio/minio/pkg/iam/policy"
	"github.com/minio/minio/pkg/madmin"
	"github.com/pkg/errors"
	"github.com/sethvargo/go-password/password"
)

type MinioUser struct {
	madmClnt *madmin.AdminClient
	logger   lager.Logger
	// preserved for subsequent connections (generating user-centric access keys)
	endpoint        string
	accessKey       string
	secretKey       string
	secure          bool
	customTransport http.RoundTripper
}

func NewMinioUser(
	logger lager.Logger,
	endpoint string,
	accessKey string,
	secretKey string,
	secure bool,
	customTransport http.RoundTripper,
) *MinioUser {
	madmClnt, err := newMinioAdminClient(endpoint, accessKey, secretKey, secure, customTransport)
	if err != nil {
		logger.Fatal("unable to create MinIO admin client", err)
	}

	return &MinioUser{
		madmClnt:        madmClnt,
		logger:          logger.Session("minio-user"),
		endpoint:        endpoint,
		accessKey:       accessKey,
		secretKey:       secretKey,
		secure:          secure,
		customTransport: customTransport,
	}
}

func newMinioAdminClient(endpoint, accessKey, secretKey string, secure bool, customTransport http.RoundTripper) (*madmin.AdminClient, error) {
	madmClnt, err := madmin.New(endpoint, accessKey, secretKey, secure)
	if err != nil {
		return nil, fmt.Errorf("Unable to connect to MinIO: %s", err)
	}
	madmClnt.SetCustomTransport(customTransport)
	return madmClnt, nil
}

func (i *MinioUser) Describe(userName string) (UserDetails, error) {
	// We simply echo data back; note that presence of the user is tested...
	userDetails := UserDetails{
		UserName: userName,
		UserARN:  toUserARN(userName),
		UserID:   userName,
	}

	i.logger.Debug("get-user", lager.Data{"input": userName})

	_, err := i.madmClnt.GetUserInfo(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return userDetails, err
	}

	i.logger.Debug("get-user", lager.Data{"output": userDetails})

	return userDetails, nil
}

func (i *MinioUser) Create(userName, iamPath string) (string, error) {
	i.logger.Debug("create-user", lager.Data{"userName": userName, "iamPath": iamPath})

	secretKey, err := password.Generate(64, 10, 0, false, true)
	if err != nil {
		i.logger.Error("password-generator", err)
		return "", err
	}

	err = i.madmClnt.AddUser(context.Background(), userName, secretKey)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", err
	}
	i.logger.Debug("create-user", lager.Data{"output": userName})

	return toUserARN(userName), nil
}

func (i *MinioUser) Delete(userName string) error {
	i.logger.Debug("delete-user", lager.Data{"userName": userName})

	err := i.madmClnt.RemoveUser(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("delete-user", lager.Data{"output": userName})

	return nil
}

func (i *MinioUser) createUserClient(userName string) (*madmin.AdminClient, error) {
	// UserInfo does not appear to respond with actual SecretKey
	// ... so we just set it again so we can auth to it and create a key!
	secretKey, err := password.Generate(64, 10, 0, false, true)
	if err != nil {
		i.logger.Error("createUserClient/password-generator", err)
		return nil, err
	}

	err = i.madmClnt.SetUser(context.Background(), userName, secretKey, madmin.AccountEnabled)
	if err != nil {
		i.logger.Error("createUserClient/minio-error", err)
		return nil, err
	}

	// Connect as alternate user
	return newMinioAdminClient(i.endpoint, userName, secretKey, i.secure, i.customTransport)
}

func (i *MinioUser) ListAccessKeys(userName string) ([]string, error) {
	i.logger.Debug("list-access-keys", lager.Data{"input": userName})

	var accessKeys []string

	// Connect as alternate user
	userClient, err := i.createUserClient(userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return accessKeys, err
	}

	accounts, err := userClient.ListServiceAccounts(context.Background())
	if err != nil {
		i.logger.Error("minio-error", err)
		return accessKeys, err
	}

	i.logger.Debug("list-access-keys", lager.Data{"output": accounts.Accounts})

	return accounts.Accounts, nil
}

func (i *MinioUser) CreateAccessKey(userName string) (string, string, error) {
	i.logger.Debug("create-access-key", lager.Data{"input": userName})

	info, err := i.madmClnt.GetUserInfo(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", "", err
	}

	var policy *iampolicy.Policy
	if info.PolicyName != "" {
		policy, err = i.madmClnt.InfoCannedPolicy(context.Background(), info.PolicyName)
		if err != nil {
			i.logger.Error("minio-error", err)
			return "", "", err
		}
	}

	// Connect as alternate user
	userClient, err := i.createUserClient(userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", "", err
	}

	creds, err := userClient.AddServiceAccount(context.Background(), policy)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", "", err
	}

	i.logger.Debug("create-access-key", lager.Data{"output": creds.AccessKey})

	return creds.AccessKey, creds.SecretKey, nil
}

func (i *MinioUser) DeleteAccessKey(userName, accessKeyID string) error {
	i.logger.Debug("delete-access-key", lager.Data{"input": userName})

	// Connect as alternate user
	userClient, err := i.createUserClient(userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	err = userClient.DeleteServiceAccount(context.Background(), accessKeyID)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("delete-access-key", lager.Data{"output": "N/A"})

	return nil
}

func (i *MinioUser) CreatePolicy(policyName, iamPath, policyTemplate string, resources []string) (string, error) {
	i.logger.Debug("create-policy", lager.Data{"policyName": policyName, "iamPath": iamPath, "policyTemplate": policyTemplate, "resources": resources})

	tmpl, err := template.New("policy").Funcs(template.FuncMap{
		"resources": func(suffix string) string {
			resourcePaths := make([]string, len(resources))
			for idx, resource := range resources {
				resourcePaths[idx] = resource + suffix
			}
			marshaled, _ := json.Marshal(resourcePaths)
			return string(marshaled)
		},
	}).Parse(policyTemplate)
	if err != nil {
		i.logger.Error("parse-error", err)
		return "", err
	}
	policy := bytes.Buffer{}
	err = tmpl.Execute(&policy, map[string]interface{}{
		"Resource":  resources[0],
		"Resources": resources,
	})
	if err != nil {
		i.logger.Error("template-error", err)
		return "", err
	}

	policyInput := iampolicy.Policy{}
	err = json.Unmarshal(policy.Bytes(), &policyInput)
	if err != nil {
		i.logger.Error("unmarshal-error", err)
		return "", err
	}
	i.logger.Debug("create-policy", lager.Data{"input": policyInput})

	err = i.madmClnt.AddCannedPolicy(context.Background(), policyName, &policyInput)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", err
	}

	i.logger.Debug("create-policy", lager.Data{"output": "created"})

	return toPolicyARN(policyName), nil
}

func (i *MinioUser) DeletePolicy(policyARN string) error {
	i.logger.Debug("delete-policy", lager.Data{"input": policyARN})

	policy, err := fromPolicyARN(policyARN)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	err = i.madmClnt.RemoveCannedPolicy(context.Background(), policy)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("delete-policy", lager.Data{"output": "removed"})

	return nil
}

func (i *MinioUser) ListAttachedUserPolicies(userName, iamPath string) ([]string, error) {
	var userPolicies []string

	i.logger.Debug("list-attached-user-policies", lager.Data{"userName": userName, "iamPath": iamPath})

	userInfo, err := i.madmClnt.GetUserInfo(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return userPolicies, err
	}
	i.logger.Debug("list-attached-user-policies", lager.Data{"output": userInfo.PolicyName})

	userPolicies = append(userPolicies, userInfo.MemberOf...)
	return userPolicies, nil
}

func (i *MinioUser) AttachUserPolicy(userName, policyARN string) error {
	i.logger.Debug("attach-user-policy", lager.Data{"userName": userName, "policyARN": policyARN})

	policy, err := fromPolicyARN(policyARN)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	err = i.madmClnt.SetPolicy(context.Background(), policy, userName, false)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("attach-user-policy", lager.Data{"output": "success"})

	return nil
}

func (i *MinioUser) DetachUserPolicy(userName, policyARN string) error {
	i.logger.Debug("detach-user-policy", lager.Data{"userName": userName, "policyARN": policyARN})

	err := i.madmClnt.SetPolicy(context.Background(), "", userName, false)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("detach-user-policy", lager.Data{"output": "success"})

	return nil
}

func toUserARN(userName string) string {
	return toARN("user", userName)
}
func toPolicyARN(policyName string) string {
	return toARN("policy", policyName)
}
func toARN(resourceType, resource string) string {
	return fmt.Sprintf("arn:aws:iam:::%s/%s", resourceType, resource)
}

func fromPolicyARN(policyARN string) (string, error) {
	return fromARN("policy", policyARN)
}
func fromARN(resourceType, theARN string) (string, error) {
	arn, err := awsarn.Parse(theARN)
	if err != nil {
		return "", nil
	}

	prefix := fmt.Sprintf("%s/", resourceType)
	if !strings.HasPrefix(arn.Resource, prefix) {
		return "", errors.Errorf("Not a %s ARN: %s", resourceType, theARN)
	}

	return strings.ReplaceAll(arn.Resource, prefix, ""), nil
}
<commit_msg>fixing policy list<commit_after>package awsiam

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"text/template"

	"code.cloudfoundry.org/lager"
	awsarn "github.com/aws/aws-sdk-go/aws/arn"
	iampolicy "github.com/minio/minio/pkg/iam/policy"
	"github.com/minio/minio/pkg/madmin"
	"github.com/pkg/errors"
	"github.com/sethvargo/go-password/password"
)

type MinioUser struct {
	madmClnt *madmin.AdminClient
	logger   lager.Logger
	// preserved for subsequent connections (generating user-centric access keys)
	endpoint        string
	accessKey       string
	secretKey       string
	secure          bool
	customTransport http.RoundTripper
}

func NewMinioUser(
	logger lager.Logger,
	endpoint string,
	accessKey string,
	secretKey string,
	secure bool,
	customTransport http.RoundTripper,
) *MinioUser {
	madmClnt, err := newMinioAdminClient(endpoint, accessKey, secretKey, secure, customTransport)
	if err != nil {
		logger.Fatal("unable to create MinIO admin client", err)
	}

	return &MinioUser{
		madmClnt:        madmClnt,
		logger:          logger.Session("minio-user"),
		endpoint:        endpoint,
		accessKey:       accessKey,
		secretKey:       secretKey,
		secure:          secure,
		customTransport: customTransport,
	}
}

func newMinioAdminClient(endpoint, accessKey, secretKey string, secure bool, customTransport http.RoundTripper) (*madmin.AdminClient, error) {
	madmClnt, err := madmin.New(endpoint, accessKey, secretKey, secure)
	if err != nil {
		return nil, fmt.Errorf("Unable to connect to MinIO: %s", err)
	}
	madmClnt.SetCustomTransport(customTransport)
	return madmClnt, nil
}

func (i *MinioUser) Describe(userName string) (UserDetails, error) {
	// We simply echo data back; note that presence of the user is tested...
	userDetails := UserDetails{
		UserName: userName,
		UserARN:  toUserARN(userName),
		UserID:   userName,
	}

	i.logger.Debug("get-user", lager.Data{"input": userName})

	_, err := i.madmClnt.GetUserInfo(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return userDetails, err
	}

	i.logger.Debug("get-user", lager.Data{"output": userDetails})

	return userDetails, nil
}

func (i *MinioUser) Create(userName, iamPath string) (string, error) {
	i.logger.Debug("create-user", lager.Data{"userName": userName, "iamPath": iamPath})

	secretKey, err := password.Generate(64, 10, 0, false, true)
	if err != nil {
		i.logger.Error("password-generator", err)
		return "", err
	}

	err = i.madmClnt.AddUser(context.Background(), userName, secretKey)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", err
	}
	i.logger.Debug("create-user", lager.Data{"output": userName})

	return toUserARN(userName), nil
}

func (i *MinioUser) Delete(userName string) error {
	i.logger.Debug("delete-user", lager.Data{"userName": userName})

	err := i.madmClnt.RemoveUser(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("delete-user", lager.Data{"output": userName})

	return nil
}

func (i *MinioUser) createUserClient(userName string) (*madmin.AdminClient, error) {
	// UserInfo does not appear to respond with actual SecretKey
	// ... so we just set it again so we can auth to it and create a key!
	secretKey, err := password.Generate(64, 10, 0, false, true)
	if err != nil {
		i.logger.Error("createUserClient/password-generator", err)
		return nil, err
	}

	err = i.madmClnt.SetUser(context.Background(), userName, secretKey, madmin.AccountEnabled)
	if err != nil {
		i.logger.Error("createUserClient/minio-error", err)
		return nil, err
	}

	// Connect as alternate user
	return newMinioAdminClient(i.endpoint, userName, secretKey, i.secure, i.customTransport)
}

func (i *MinioUser) ListAccessKeys(userName string) ([]string, error) {
	i.logger.Debug("list-access-keys", lager.Data{"input": userName})

	var accessKeys []string

	// Connect as alternate user
	userClient, err := i.createUserClient(userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return accessKeys, err
	}

	accounts, err := userClient.ListServiceAccounts(context.Background())
	if err != nil {
		i.logger.Error("minio-error", err)
		return accessKeys, err
	}

	i.logger.Debug("list-access-keys", lager.Data{"output": accounts.Accounts})

	return accounts.Accounts, nil
}

func (i *MinioUser) CreateAccessKey(userName string) (string, string, error) {
	i.logger.Debug("create-access-key", lager.Data{"input": userName})

	info, err := i.madmClnt.GetUserInfo(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", "", err
	}

	var policy *iampolicy.Policy
	if info.PolicyName != "" {
		policy, err = i.madmClnt.InfoCannedPolicy(context.Background(), info.PolicyName)
		if err != nil {
			i.logger.Error("minio-error", err)
			return "", "", err
		}
	}

	// Connect as alternate user
	userClient, err := i.createUserClient(userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", "", err
	}

	creds, err := userClient.AddServiceAccount(context.Background(), policy)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", "", err
	}

	i.logger.Debug("create-access-key", lager.Data{"output": creds.AccessKey})

	return creds.AccessKey, creds.SecretKey, nil
}

func (i *MinioUser) DeleteAccessKey(userName, accessKeyID string) error {
	i.logger.Debug("delete-access-key", lager.Data{"input": userName})

	// Connect as alternate user
	userClient, err := i.createUserClient(userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	err = userClient.DeleteServiceAccount(context.Background(), accessKeyID)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("delete-access-key", lager.Data{"output": "N/A"})

	return nil
}

func (i *MinioUser) CreatePolicy(policyName, iamPath, policyTemplate string, resources []string) (string, error) {
	i.logger.Debug("create-policy", lager.Data{"policyName": policyName, "iamPath": iamPath, "policyTemplate": policyTemplate, "resources": resources})

	tmpl, err := template.New("policy").Funcs(template.FuncMap{
		"resources": func(suffix string) string {
			resourcePaths := make([]string, len(resources))
			for idx, resource := range resources {
				resourcePaths[idx] = resource + suffix
			}
			marshaled, _ := json.Marshal(resourcePaths)
			return string(marshaled)
		},
	}).Parse(policyTemplate)
	if err != nil {
		i.logger.Error("parse-error", err)
		return "", err
	}
	policy := bytes.Buffer{}
	err = tmpl.Execute(&policy, map[string]interface{}{
		"Resource":  resources[0],
		"Resources": resources,
	})
	if err != nil {
		i.logger.Error("template-error", err)
		return "", err
	}

	policyInput := iampolicy.Policy{}
	err = json.Unmarshal(policy.Bytes(), &policyInput)
	if err != nil {
		i.logger.Error("unmarshal-error", err)
		return "", err
	}
	i.logger.Debug("create-policy", lager.Data{"input": policyInput})

	err = i.madmClnt.AddCannedPolicy(context.Background(), policyName, &policyInput)
	if err != nil {
		i.logger.Error("minio-error", err)
		return "", err
	}

	i.logger.Debug("create-policy", lager.Data{"output": "created"})

	return toPolicyARN(policyName), nil
}

func (i *MinioUser) DeletePolicy(policyARN string) error {
	i.logger.Debug("delete-policy", lager.Data{"input": policyARN})

	policy, err := fromPolicyARN(policyARN)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	err = i.madmClnt.RemoveCannedPolicy(context.Background(), policy)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("delete-policy", lager.Data{"output": "removed"})

	return nil
}

func (i *MinioUser) ListAttachedUserPolicies(userName, iamPath string) ([]string, error) {
	var userPolicies []string

	i.logger.Debug("list-attached-user-policies", lager.Data{"userName": userName, "iamPath": iamPath})

	userInfo, err := i.madmClnt.GetUserInfo(context.Background(), userName)
	if err != nil {
		i.logger.Error("minio-error", err)
		return userPolicies, err
	}
	i.logger.Debug("list-attached-user-policies", lager.Data{"output": userInfo.PolicyName})

	userPolicies = append(userPolicies, userInfo.PolicyName)
	return userPolicies, nil
}

func (i *MinioUser) AttachUserPolicy(userName, policyARN string) error {
	i.logger.Debug("attach-user-policy", lager.Data{"userName": userName, "policyARN": policyARN})

	policy, err := fromPolicyARN(policyARN)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	err = i.madmClnt.SetPolicy(context.Background(), policy, userName, false)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("attach-user-policy", lager.Data{"output": "success"})

	return nil
}

func (i *MinioUser) DetachUserPolicy(userName, policyARN string) error {
	i.logger.Debug("detach-user-policy", lager.Data{"userName": userName, "policyARN": policyARN})

	err := i.madmClnt.SetPolicy(context.Background(), "", userName, false)
	if err != nil {
		i.logger.Error("minio-error", err)
		return err
	}

	i.logger.Debug("detach-user-policy", lager.Data{"output": "success"})

	return nil
}

func toUserARN(userName string) string {
	return toARN("user", userName)
}
func toPolicyARN(policyName string) string {
	return toARN("policy", policyName)
}
func toARN(resourceType, resource string) string {
	return fmt.Sprintf("arn:aws:iam:::%s/%s", resourceType, resource)
}

func fromPolicyARN(policyARN string) (string, error) {
	return fromARN("policy", policyARN)
}
func fromARN(resourceType, theARN string) (string, error) {
	arn, err := awsarn.Parse(theARN)
	if err != nil {
		return "", nil
	}

	prefix := fmt.Sprintf("%s/", resourceType)
	if !strings.HasPrefix(arn.Resource, prefix) {
		return "", errors.Errorf("Not a %s ARN: %s", resourceType, theARN)
	}

	return strings.ReplaceAll(arn.Resource, prefix, ""), nil
}
<|endoftext|>


<commit_before>package httpexpect

import (
	"testing"

	"github.com/gorilla/websocket"
)

func TestWebsocketFailed(t *testing.T) {
	chain := makeChain(newMockReporter(t))

	chain.fail("fail")

	ws := &Websocket{
		chain: chain,
	}

	ws.chain.assertFailed(t)

	ws.Raw()
	ws.WithReadTimeout(0)
	ws.WithoutReadTimeout()
	ws.WithWriteTimeout(0)
	ws.WithoutWriteTimeout()

	ws.Subprotocol().chain.assertFailed(t)
	ws.Expect().chain.assertFailed(t)

	ws.WriteMessage(websocket.TextMessage, []byte("a"))
	ws.WriteBytesBinary([]byte("a"))
	ws.WriteBytesText([]byte("a"))
	ws.WriteText("a")
	ws.WriteJSON(map[string]string{"a": "b"})

	ws.Close()
	ws.CloseWithBytes([]byte("a"))
	ws.CloseWithJSON(map[string]string{"a": "b"})
	ws.CloseWithText("a")

	ws.Disconnect()
}

func TestWebsocketNil(t *testing.T) {
	config := Config{
		Reporter: newMockReporter(t),
	}

	ws := NewWebsocket(config, nil)

	msg := ws.Expect()
	msg.chain.assertFailed(t)

	ws.chain.assertFailed(t)
}

func TestWebsocket_Expect(t *testing.T) {

	t.Run("websocket is closed", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: true,
		}

		c.Expect()

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected read from closed WebSocket connection"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}

func TestWebsocket_checkUnusable(t *testing.T) {
	type fields struct {
		conn     *websocket.Conn
		isClosed bool
	}
	type args struct {
		reporter *stringReporter
		where    string
	}
	tests := []struct {
		name     string
		fields   fields
		args     args
		want     bool
		reported bool
		wantMsg  string
	}{
		{
			name: "conn is nil",
			fields: fields{
				conn:     nil,
				isClosed: false,
			},
			args: args{
				reporter: newStringReporter(),
				where:    "Close",
			},
			want:     true,
			reported: true,
			wantMsg:  "\nunexpected Close call for failed WebSocket connection",
		},
		{
			name: "websocket is closed",
			fields: fields{
				conn:     &websocket.Conn{},
				isClosed: true,
			},
			args: args{
				reporter: newStringReporter(),
				where:    "Close",
			},
			want:     true,
			reported: true,
			wantMsg:  "\nunexpected Close call for closed WebSocket connection",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Websocket{
				chain:    makeChain(tt.args.reporter),
				conn:     tt.fields.conn,
				isClosed: tt.fields.isClosed,
			}
			if got := c.checkUnusable(tt.args.where); got != tt.want {
				t.Errorf("Websocket.checkUnusable() = %v, want %v", got, tt.want)
				return
			}
			if got := tt.args.reporter.reported; got != tt.reported {
				t.Errorf("Websocket.checkUnusable() is error message reported = %v, want %v", got, tt.want)
				return
			}
			if got := tt.args.reporter.msg; got != tt.wantMsg {
				t.Errorf("Websocket.checkUnusable() error message = %v, want %v", got, tt.want)
				return
			}
		})
	}
}

func TestWriteJSONMarshalFail(t *testing.T) {
	r := newStringReporter()

	ws := &Websocket{
		chain:    makeChain(r),
		conn:     &websocket.Conn{},
		isClosed: false,
	}

	channel := make(chan int)

	ws.WriteJSON(channel)

	ws.chain.assertFailed(t)

	if !r.reported {
		t.Errorf("Error message not reported")
		return
	}
}

func TestWebsocket_WriteMessage(t *testing.T) {
	type args struct {
		reporter  *stringReporter
		typ       int
		content   []byte
		closeCode []int
	}
	tests := []struct {
		name        string
		args        args
		reported    bool
		reportedMsg string
	}{
		{
			name: "Close message, multiple close code",
			args: args{
				reporter:  newStringReporter(),
				typ:       websocket.CloseMessage,
				content:   []byte("closing message..."),
				closeCode: []int{websocket.CloseNormalClosure, websocket.CloseAbnormalClosure},
			},
			reported:    true,
			reportedMsg: "\nunexpected multiple closeCode arguments passed to WriteMessage",
		},
		{
			name: "Close message, multiple close code",
			args: args{
				reporter:  newStringReporter(),
				typ:       websocket.PingMessage,
				content:   []byte("closing message..."),
				closeCode: []int{websocket.CloseNormalClosure, websocket.CloseAbnormalClosure},
			},
			reported:    true,
			reportedMsg: "\nunexpected WebSocket message type 'ping' passed to WriteMessage",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Websocket{
				chain:    makeChain(tt.args.reporter),
				conn:     &websocket.Conn{},
				isClosed: false,
			}

			c.WriteMessage(tt.args.typ, tt.args.content, tt.args.closeCode...)

			if got := tt.args.reporter.reported; got != tt.reported {
				t.Errorf("Websocket.WriteMessage() is error message reported = %v, want %v", got, tt.reported)
			}

			if got := tt.args.reporter.msg; got != tt.reportedMsg {
				t.Errorf("Websocket.WriteMessage() error message = %v, want %v", got, tt.reportedMsg)
			}
		})
	}
}

func TestWebsocket_CloseWithText(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithText("Closing...", websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to CloseWithText"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}

func TestWebsocket_CloseWithJSON(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithJSON("Closing...", websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to CloseWithJSON"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})

	t.Run("json marshall error", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithJSON(make(chan int), websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}
	})
}

func TestWebsocket_CloseWithBytes(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithBytes([]byte("Closing..."), websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to CloseWithBytes"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}

func TestWebsocket_Close(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.Close(websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to Close"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}
<commit_msg>Remove underscores from new unit test's names<commit_after>package httpexpect

import (
	"testing"

	"github.com/gorilla/websocket"
)

func TestWebsocketFailed(t *testing.T) {
	chain := makeChain(newMockReporter(t))

	chain.fail("fail")

	ws := &Websocket{
		chain: chain,
	}

	ws.chain.assertFailed(t)

	ws.Raw()
	ws.WithReadTimeout(0)
	ws.WithoutReadTimeout()
	ws.WithWriteTimeout(0)
	ws.WithoutWriteTimeout()

	ws.Subprotocol().chain.assertFailed(t)
	ws.Expect().chain.assertFailed(t)

	ws.WriteMessage(websocket.TextMessage, []byte("a"))
	ws.WriteBytesBinary([]byte("a"))
	ws.WriteBytesText([]byte("a"))
	ws.WriteText("a")
	ws.WriteJSON(map[string]string{"a": "b"})

	ws.Close()
	ws.CloseWithBytes([]byte("a"))
	ws.CloseWithJSON(map[string]string{"a": "b"})
	ws.CloseWithText("a")

	ws.Disconnect()
}

func TestWebsocketNil(t *testing.T) {
	config := Config{
		Reporter: newMockReporter(t),
	}

	ws := NewWebsocket(config, nil)

	msg := ws.Expect()
	msg.chain.assertFailed(t)

	ws.chain.assertFailed(t)
}

func TestWebsocketExpect(t *testing.T) {

	t.Run("websocket is closed", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: true,
		}

		c.Expect()

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected read from closed WebSocket connection"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}

func TestWebsocketCheckUnusable(t *testing.T) {
	type fields struct {
		conn     *websocket.Conn
		isClosed bool
	}
	type args struct {
		reporter *stringReporter
		where    string
	}
	tests := []struct {
		name     string
		fields   fields
		args     args
		want     bool
		reported bool
		wantMsg  string
	}{
		{
			name: "conn is nil",
			fields: fields{
				conn:     nil,
				isClosed: false,
			},
			args: args{
				reporter: newStringReporter(),
				where:    "Close",
			},
			want:     true,
			reported: true,
			wantMsg:  "\nunexpected Close call for failed WebSocket connection",
		},
		{
			name: "websocket is closed",
			fields: fields{
				conn:     &websocket.Conn{},
				isClosed: true,
			},
			args: args{
				reporter: newStringReporter(),
				where:    "Close",
			},
			want:     true,
			reported: true,
			wantMsg:  "\nunexpected Close call for closed WebSocket connection",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Websocket{
				chain:    makeChain(tt.args.reporter),
				conn:     tt.fields.conn,
				isClosed: tt.fields.isClosed,
			}
			if got := c.checkUnusable(tt.args.where); got != tt.want {
				t.Errorf("Websocket.checkUnusable() = %v, want %v", got, tt.want)
				return
			}
			if got := tt.args.reporter.reported; got != tt.reported {
				t.Errorf("Websocket.checkUnusable() is error message reported = %v, want %v", got, tt.want)
				return
			}
			if got := tt.args.reporter.msg; got != tt.wantMsg {
				t.Errorf("Websocket.checkUnusable() error message = %v, want %v", got, tt.want)
				return
			}
		})
	}
}

func TestWebsocketWriteJSONMarshalFail(t *testing.T) {
	r := newStringReporter()

	ws := &Websocket{
		chain:    makeChain(r),
		conn:     &websocket.Conn{},
		isClosed: false,
	}

	channel := make(chan int)

	ws.WriteJSON(channel)

	ws.chain.assertFailed(t)

	if !r.reported {
		t.Errorf("Error message not reported")
		return
	}
}

func TestWebsocketWriteMessage(t *testing.T) {
	type args struct {
		reporter  *stringReporter
		typ       int
		content   []byte
		closeCode []int
	}
	tests := []struct {
		name        string
		args        args
		reported    bool
		reportedMsg string
	}{
		{
			name: "Close message, multiple close code",
			args: args{
				reporter:  newStringReporter(),
				typ:       websocket.CloseMessage,
				content:   []byte("closing message..."),
				closeCode: []int{websocket.CloseNormalClosure, websocket.CloseAbnormalClosure},
			},
			reported:    true,
			reportedMsg: "\nunexpected multiple closeCode arguments passed to WriteMessage",
		},
		{
			name: "Close message, multiple close code",
			args: args{
				reporter:  newStringReporter(),
				typ:       websocket.PingMessage,
				content:   []byte("closing message..."),
				closeCode: []int{websocket.CloseNormalClosure, websocket.CloseAbnormalClosure},
			},
			reported:    true,
			reportedMsg: "\nunexpected WebSocket message type 'ping' passed to WriteMessage",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Websocket{
				chain:    makeChain(tt.args.reporter),
				conn:     &websocket.Conn{},
				isClosed: false,
			}

			c.WriteMessage(tt.args.typ, tt.args.content, tt.args.closeCode...)

			if got := tt.args.reporter.reported; got != tt.reported {
				t.Errorf("Websocket.WriteMessage() is error message reported = %v, want %v", got, tt.reported)
			}

			if got := tt.args.reporter.msg; got != tt.reportedMsg {
				t.Errorf("Websocket.WriteMessage() error message = %v, want %v", got, tt.reportedMsg)
			}
		})
	}
}

func TestWebsocketCloseWithText(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithText("Closing...", websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to CloseWithText"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}

func TestWebsocketCloseWithJSON(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithJSON("Closing...", websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to CloseWithJSON"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})

	t.Run("json marshall error", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithJSON(make(chan int), websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}
	})
}

func TestWebsocketCloseWithBytes(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.CloseWithBytes([]byte("Closing..."), websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to CloseWithBytes"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}

func TestWebsocketClose(t *testing.T) {
	t.Run("multiple code args", func(t *testing.T) {
		r := newStringReporter()
		c := &Websocket{
			chain:    makeChain(r),
			conn:     &websocket.Conn{},
			isClosed: false,
		}

		c.Close(websocket.CloseNormalClosure, websocket.CloseAbnormalClosure)

		if !r.reported {
			t.Errorf("Websocket.CloseWithText() error message not reported")
		}

		want := "\nunexpected multiple code arguments passed to Close"

		if got := r.msg; got != want {
			t.Errorf("Websocket.CloseWithText() = %v, want %v", got, want)
		}
	})
}
<|endoftext|>


<commit_before>package wsdlgen

import "aqwari.net/xml/internal/gen"

// One of the goals of this package is that generated code
// has no external dependencies, only the Go standard
// library. That means we have to bundle any static
// "helper" functions along with the generated code. We
// are playing a balancing game here; the larger the static
// code base grows, the weaker the argument against external
// dependencies becomes.
var helpers string = `
	type Client struct {
		HTTPClient http.Client
		ResponseHook func(*http.Response)
		RequestHook func(*http.Request)
	}

	type soapEnvelope struct {
		XMLName struct{} ` + "`" + `xml:"http://schemas.xmlsoap.org/soap/envelope/ Envelope"` + "`" + `
		Header []byte ` + "`" + `xml:"http://schemas.xmlsoap.org/soap/envelope/ Header"` + "`" + `
		Body struct {
			Message interface{}
			Fault struct {
				String string ` + "`xml:\"faultstring\"`" + `
				Code string ` + "`xml:\"faultcode\"`" + `
				Detail string ` + "`xml:\"detail\"`" + `
			} ` + "`xml:\"http://schemas.xmlsoap.org/soap/envelope/ Fault\"`" + `
		}` + "`" + `xml:"http://schemas.xmlsoap.org/soap/envelope/ Body"` + "`" + `
	}

	func (c *Client) do(method, uri, action string, in, out interface{}) error {
		var body io.Reader
		var envelope soapEnvelope

		if method == "POST" || method == "PUT" {
			var buf bytes.Buffer
			envelope.Body.Message = in
			enc := xml.NewEncoder(&buf)
			if err := enc.Encode(envelope); err != nil {
				return err
			}
			if err := enc.Flush(); err != nil {
				return err
			}
			body = &buf
		}
		req, err := http.NewRequest(method, uri, body)
		if err != nil {
			return err
		}
		req.Header.Set("SOAPAction", action)
		if c.RequestHook != nil {
			c.RequestHook(req)
		}
		rsp, err := c.HTTPClient.Do(req)
		if err != nil {
			return err
		}
		defer rsp.Body.Close()

		if c.ResponseHook != nil {
			c.ResponseHook(rsp)
		}

		dec := xml.NewDecoder(rsp.Body)
		envelope.Body.Message = out
		if err := dec.Decode(&envelope); err != nil {
			return err
		}
		if envelope.Body.Fault.Code != "" || envelope.Body.Fault.String != "" {
			return fmt.Errorf("%s: %s", envelope.Body.Fault.Code, envelope.Body.Fault.String)
		}
		return nil
	}
`

func (p *printer) addHelpers() {
	decls, err := gen.Declarations(helpers)
	if err != nil {
		// code does not change at runtime, so
		// this should never happen
		panic(err)
	}
	p.file.Decls = append(p.file.Decls, decls...)
}
<commit_msg>Do not send/receive soap Fault if it is empty<commit_after>package wsdlgen

import "aqwari.net/xml/internal/gen"

// One of the goals of this package is that generated code
// has no external dependencies, only the Go standard
// library. That means we have to bundle any static
// "helper" functions along with the generated code. We
// are playing a balancing game here; the larger the static
// code base grows, the weaker the argument against external
// dependencies becomes.
var helpers string = `
	type Client struct {
		HTTPClient http.Client
		ResponseHook func(*http.Response)
		RequestHook func(*http.Request)
	}

	type soapEnvelope struct {
		XMLName struct{} ` + "`" + `xml:"http://schemas.xmlsoap.org/soap/envelope/ Envelope"` + "`" + `
		Header []byte ` + "`" + `xml:"http://schemas.xmlsoap.org/soap/envelope/ Header"` + "`" + `
		Body struct {
			Message interface{}
			Fault *struct {
				String string ` + "`xml:\"faultstring,omitempty\"`" + `
				Code string ` + "`xml:\"faultcode,omitempty\"`" + `
				Detail string ` + "`xml:\"detail,omitempty\"`" + `
			} ` + "`xml:\"http://schemas.xmlsoap.org/soap/envelope/ Fault,omitempty\"`" + `
		}` + "`" + `xml:"http://schemas.xmlsoap.org/soap/envelope/ Body"` + "`" + `
	}

	func (c *Client) do(method, uri, action string, in, out interface{}) error {
		var body io.Reader
		var envelope soapEnvelope

		if method == "POST" || method == "PUT" {
			var buf bytes.Buffer
			envelope.Body.Message = in
			enc := xml.NewEncoder(&buf)
			if err := enc.Encode(envelope); err != nil {
				return err
			}
			if err := enc.Flush(); err != nil {
				return err
			}
			body = &buf
		}
		req, err := http.NewRequest(method, uri, body)
		if err != nil {
			return err
		}
		req.Header.Set("SOAPAction", action)
		if c.RequestHook != nil {
			c.RequestHook(req)
		}
		rsp, err := c.HTTPClient.Do(req)
		if err != nil {
			return err
		}
		defer rsp.Body.Close()

		if c.ResponseHook != nil {
			c.ResponseHook(rsp)
		}

		dec := xml.NewDecoder(rsp.Body)
		envelope.Body.Message = out
		if err := dec.Decode(&envelope); err != nil {
			return err
		}
		if envelope.Body.Fault != nil {
			return fmt.Errorf("%s: %s", envelope.Body.Fault.Code, envelope.Body.Fault.String)
		}
		return nil
	}
`

func (p *printer) addHelpers() {
	decls, err := gen.Declarations(helpers)
	if err != nil {
		// code does not change at runtime, so
		// this should never happen
		panic(err)
	}
	p.file.Decls = append(p.file.Decls, decls...)
}
<|endoftext|>


<commit_before>package xslt

import (
	"fmt"
	"github.com/moovweb/gokogiri/xml"
	"github.com/moovweb/gokogiri/xpath"
	"unsafe"
)

func (style *Stylesheet) RegisterXsltFunctions() {
	style.Functions["key"] = XsltKey
	style.Functions["system-property"] = XsltSystemProperty
	style.Functions["document"] = XsltDocumentFn
	//element-available
	//function-available
	//document
	//id
	//current
	//lang
	//generate-id
	//unparsed-entity-uri
	//format-number
}

type Nodeset []xml.Node

type Key struct {
	nodes map[string]Nodeset
	use   string
	match string
}

/*
func (key *Key) Evaluate() {
    c := CompileMatch()
    for n in doc.Nodes {
        if c.Matches(n) {
            Nodes = append(Nodes, n)
        }
    }
}
*/

func (n Nodeset) ToPointers() (pointers []unsafe.Pointer) {
	for _, node := range n {
		pointers = append(pointers, node.NodePtr())
	}
	return
}

// Implementation of key() from XSLT spec
func XsltKey(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 2 {
		return nil
	}
	// always convert to string
	name := args[0].(string)
	// convert to string (TODO: unless nodeset)
	val := args[1].(string)
	//get the execution context
	c := context.(*ExecutionContext)
	//look up the key
	k, ok := c.Style.Keys[name]
	if !ok {
		return nil
	}
	result, _ := k.nodes[val]
	//return the nodeset
	return result.ToPointers()
}

// Implementation of system-property() from XSLT spec
func XsltSystemProperty(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	switch args[0].(string) {
	case "xsl:version":
		return 1.0
	case "xsl:vendor":
		return "John C Barstow"
	case "xsl:vendor-url":
		return "http://github.com/jbowtie/ratago"
	default:
		fmt.Println("EXEC system-property", args[0])
	}
	return nil
}

//Implementation of document() from XSLT spec
func XsltDocumentFn(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)

	switch doc := args[0].(type) {
	case string:
		if doc == "" {
			return []xml.Node{c.Style.Doc}
		}
		return nil
	}
	return nil
}
<commit_msg>Make note that a nodeset can be passed as second argument to xsl:key<commit_after>package xslt

import (
	"fmt"
	"github.com/moovweb/gokogiri/xml"
	"github.com/moovweb/gokogiri/xpath"
	"unsafe"
)

func (style *Stylesheet) RegisterXsltFunctions() {
	style.Functions["key"] = XsltKey
	style.Functions["system-property"] = XsltSystemProperty
	style.Functions["document"] = XsltDocumentFn
	//element-available
	//function-available
	//document
	//id
	//current
	//lang
	//generate-id
	//unparsed-entity-uri
	//format-number
}

type Nodeset []xml.Node

type Key struct {
	nodes map[string]Nodeset
	use   string
	match string
}

/*
func (key *Key) Evaluate() {
    c := CompileMatch()
    for n in doc.Nodes {
        if c.Matches(n) {
            Nodes = append(Nodes, n)
        }
    }
}
*/

func (n Nodeset) ToPointers() (pointers []unsafe.Pointer) {
	for _, node := range n {
		pointers = append(pointers, node.NodePtr())
	}
	return
}

// Implementation of key() from XSLT spec
func XsltKey(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 2 {
		return nil
	}
	// always convert to string
	name := args[0].(string)
	// convert to string (TODO: unless nodeset)
	val := ""
	switch v := args[1].(type) {
	case string:
		val = v
	case []unsafe.Pointer:
		// nodeset; see xsl:key spec for how to handle this
	}
	//get the execution context
	c := context.(*ExecutionContext)
	//look up the key
	k, ok := c.Style.Keys[name]
	if !ok {
		return nil
	}
	result, _ := k.nodes[val]
	//return the nodeset
	return result.ToPointers()
}

// Implementation of system-property() from XSLT spec
func XsltSystemProperty(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	switch args[0].(string) {
	case "xsl:version":
		return 1.0
	case "xsl:vendor":
		return "John C Barstow"
	case "xsl:vendor-url":
		return "http://github.com/jbowtie/ratago"
	default:
		fmt.Println("EXEC system-property", args[0])
	}
	return nil
}

//Implementation of document() from XSLT spec
func XsltDocumentFn(context xpath.VariableScope, args []interface{}) interface{} {
	if len(args) < 1 {
		return nil
	}
	c := context.(*ExecutionContext)

	switch doc := args[0].(type) {
	case string:
		if doc == "" {
			return []xml.Node{c.Style.Doc}
		}
		return nil
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings

// Count UTF-8 sequences in s.
// Assumes s is well-formed.
export func utflen(s string) int {
	n := 0;
	for i := 0; i < len(s); i++ {
		if s[i]&0xC0 != 0x80 {
			n++
		}
	}
	return n
}

// Split string into array of UTF-8 sequences (still strings)
export func explode(s string) *[]string {
	a := new([]string, utflen(s));
	j := 0;
	for i := 0; i < len(a); i++ {
		ej := j;
		ej++;
		for ej < len(s) && (s[ej]&0xC0) == 0x80 {
			ej++
		}
		a[i] = s[j:ej];
		j = ej
	}
	return a
}

// Count non-overlapping instances of sep in s.
export func count(s, sep string) int {
	if sep == "" {
		return utflen(s)+1
	}
	c := sep[0];
	n := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			n++;
			i += len(sep)-1
		}
	}
	return n
}

// Return index of first instance of sep in s.
export func index(s, sep string) int {
	if sep == "" {
		return 0
	}
	c := sep[0];
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			return i
		}
	}
	return -1
}

// Split string into list of strings at separators
export func split(s, sep string) *[]string {
	if sep == "" {
		return explode(s)
	}
	c := sep[0];
	start := 0;
	n := count(s, sep)+1;
	a := new([]string, n);
	na := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			a[na] = s[start:i];
			na++;
			start = i+len(sep);
			i += len(sep)-1
		}
	}
	a[na] = s[start:len(s)];
	return a
}

// Join list of strings with separators between them.
export func join(a *[]string, sep string) string {
	if len(a) == 0 {
		return ""
	}
	if len(a) == 1 {
		return a[0]
	}
	n := len(sep) * (len(a)-1);
	for i := 0; i < len(a); i++ {
		n += len(a[i])
	}

	b := new([]byte, n);
	bp := 0;
	for i := 0; i < len(a); i++ {
		s := a[i];
		for j := 0; j < len(s); j++ {
			b[bp] = s[j];
			bp++
		}
		if i + 1 < len(a) {
			s = sep;
			for j := 0; j < len(s); j++ {
				b[bp] = s[j];
				bp++
			}
		}
	}
	return string(b)
}

// Convert decimal string to integer.
// TODO: Doesn't check for overflow.
export func atol(s string) (i int64, ok bool) {
	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off leading sign
	neg := false;
	if s[0] == '+' {
		s = s[1:len(s)]
	} else if s[0] == '-' {
		neg = true;
		s = s[1:len(s)]
	}

	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off zero
	if s == "0" {
		return 0, true
	}

	// otherwise, leading zero bad
	if s[0] == '0' {
		return 0, false
	}

	// parse number
	n := int64(0);
	for i := 0; i < len(s); i++ {
		if s[i] < '0' || s[i] > '9' {
			return 0, false
		}
		n = n*10 + int64(s[i] - '0')
	}
	if neg {
		n = -n
	}
	return n, true
}

export func atoi(s string) (i int, ok bool) {
	ii, okok := atol(s);
	i = int(ii);
	return i, okok
}

export func ltoa(i int64) string {
	if i == 0 {
		return "0"
	}

	neg := false;	// negative
	u := uint(i);
	if i < 0 {
		neg = true;
		u = -u;
	}

	// Assemble decimal in reverse order.
	var b [32]byte;
	bp := len(b);
	for ; u > 0; u /= 10 {
		bp--;
		b[bp] = byte(u%10) + '0'
	}
	if neg {	// add sign
		bp--;
		b[bp] = '-'
	}

	// BUG return string(b[bp:len(b)])
	return string((&b)[bp:len(b)])
}

export func itoa(i int) string {
	return ltoa(int64(i));
}

// Convert float64 to string.  No control over format.
// Result not great; only useful for simple debugging.
export func dtoa(v float64) string {
	var buf [20]byte;

	const n = 7;	// digits printed
	e := 0;	// exp
	var sign byte = '+';
	if(v != 0) {
		// sign
		if(v < 0) {
			v = -v;
			sign = '-';
		}

		// normalize
		for v >= 10 {
			e++;
			v /= 10;
		}
		for v < 1 {
			e--;
			v *= 10;
		}

		// round
		var h float64 = 5;
		for i := 0; i < n; i++ {
			h /= 10;
		}
		v += h;
		if v >= 10 {
			e++;
			v /= 10;
		}
	}

	// format +d.dddd+edd
	buf[0] = sign;
	for i := 0; i < n; i++ {
		s := int64(v);
		buf[i+2] = byte(s)+'0';
		v -= float64(s);
		v *= 10;
	}
	buf[1] = buf[2];
	buf[2] = '.';

	buf[n+2] = 'e';
	buf[n+3] = '+';
	if e < 0 {
		e = -e;
		buf[n+3] = '-';
	}

	// TODO: exponents > 99?
	buf[n+4] = byte((e/10) + '0');
	buf[n+5] = byte((e%10) + '0');
	return string(buf)[0:n+6];	// TODO: should be able to slice buf
}

export func ftoa(v float) string {
	return dtoa(float64(v));
}
<commit_msg>ascii to unsigned integer<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings

// Count UTF-8 sequences in s.
// Assumes s is well-formed.
export func utflen(s string) int {
	n := 0;
	for i := 0; i < len(s); i++ {
		if s[i]&0xC0 != 0x80 {
			n++
		}
	}
	return n
}

// Split string into array of UTF-8 sequences (still strings)
export func explode(s string) *[]string {
	a := new([]string, utflen(s));
	j := 0;
	for i := 0; i < len(a); i++ {
		ej := j;
		ej++;
		for ej < len(s) && (s[ej]&0xC0) == 0x80 {
			ej++
		}
		a[i] = s[j:ej];
		j = ej
	}
	return a
}

// Count non-overlapping instances of sep in s.
export func count(s, sep string) int {
	if sep == "" {
		return utflen(s)+1
	}
	c := sep[0];
	n := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			n++;
			i += len(sep)-1
		}
	}
	return n
}

// Return index of first instance of sep in s.
export func index(s, sep string) int {
	if sep == "" {
		return 0
	}
	c := sep[0];
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			return i
		}
	}
	return -1
}

// Split string into list of strings at separators
export func split(s, sep string) *[]string {
	if sep == "" {
		return explode(s)
	}
	c := sep[0];
	start := 0;
	n := count(s, sep)+1;
	a := new([]string, n);
	na := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			a[na] = s[start:i];
			na++;
			start = i+len(sep);
			i += len(sep)-1
		}
	}
	a[na] = s[start:len(s)];
	return a
}

// Join list of strings with separators between them.
export func join(a *[]string, sep string) string {
	if len(a) == 0 {
		return ""
	}
	if len(a) == 1 {
		return a[0]
	}
	n := len(sep) * (len(a)-1);
	for i := 0; i < len(a); i++ {
		n += len(a[i])
	}

	b := new([]byte, n);
	bp := 0;
	for i := 0; i < len(a); i++ {
		s := a[i];
		for j := 0; j < len(s); j++ {
			b[bp] = s[j];
			bp++
		}
		if i + 1 < len(a) {
			s = sep;
			for j := 0; j < len(s); j++ {
				b[bp] = s[j];
				bp++
			}
		}
	}
	return string(b)
}

// Convert decimal string to unsigned integer.
// TODO: Doesn't check for overflow.
export func atoui64(s string) (i uint64, ok bool) {
	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off zero
	if s == "0" {
		return 0, true
	}

	// otherwise, leading zero bad
	if s[0] == '0' {
		return 0, false
	}

	// parse number
	n := uint64(0);
	for i := 0; i < len(s); i++ {
		if s[i] < '0' || s[i] > '9' {
			return 0, false
		}
		n = n*10 + uint64(s[i] - '0')
	}
	return n, true
}

// Convert decimal string to integer.
// TODO: Doesn't check for overflow.
export func atoi64(s string) (i int64, ok bool) {
	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off leading sign
	neg := false;
	if s[0] == '+' {
		s = s[1:len(s)]
	} else if s[0] == '-' {
		neg = true;
		s = s[1:len(s)]
	}

	var un uint64;
	un, ok = atoui64(s);
	if !ok {
		return 0, false
	}
	n := int64(un);
	if neg {
		n = -n
	}
	return n, true
}

export func atoui(s string) (i uint, ok bool) {
	ii, okok := atoui64(s);
	i = uint(ii);
	return i, okok
}

export func atoi(s string) (i int, ok bool) {
	ii, okok := atoi64(s);
	i = int(ii);
	return i, okok
}

export func ltoa(i int64) string {
	if i == 0 {
		return "0"
	}

	neg := false;	// negative
	u := uint(i);
	if i < 0 {
		neg = true;
		u = -u;
	}

	// Assemble decimal in reverse order.
	var b [32]byte;
	bp := len(b);
	for ; u > 0; u /= 10 {
		bp--;
		b[bp] = byte(u%10) + '0'
	}
	if neg {	// add sign
		bp--;
		b[bp] = '-'
	}

	// BUG return string(b[bp:len(b)])
	return string((&b)[bp:len(b)])
}

export func itoa(i int) string {
	return ltoa(int64(i));
}

// Convert float64 to string.  No control over format.
// Result not great; only useful for simple debugging.
export func dtoa(v float64) string {
	var buf [20]byte;

	const n = 7;	// digits printed
	e := 0;	// exp
	var sign byte = '+';
	if(v != 0) {
		// sign
		if(v < 0) {
			v = -v;
			sign = '-';
		}

		// normalize
		for v >= 10 {
			e++;
			v /= 10;
		}
		for v < 1 {
			e--;
			v *= 10;
		}

		// round
		var h float64 = 5;
		for i := 0; i < n; i++ {
			h /= 10;
		}
		v += h;
		if v >= 10 {
			e++;
			v /= 10;
		}
	}

	// format +d.dddd+edd
	buf[0] = sign;
	for i := 0; i < n; i++ {
		s := int64(v);
		buf[i+2] = byte(s)+'0';
		v -= float64(s);
		v *= 10;
	}
	buf[1] = buf[2];
	buf[2] = '.';

	buf[n+2] = 'e';
	buf[n+3] = '+';
	if e < 0 {
		e = -e;
		buf[n+3] = '-';
	}

	// TODO: exponents > 99?
	buf[n+4] = byte((e/10) + '0');
	buf[n+5] = byte((e%10) + '0');
	return string(buf)[0:n+6];	// TODO: should be able to slice buf
}

export func ftoa(v float) string {
	return dtoa(float64(v));
}
<|endoftext|>


<commit_before>package gubled

import (
	"github.com/smancke/guble/client"
	"github.com/smancke/guble/gcm"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/testutil"

	"github.com/stretchr/testify/assert"

	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"testing"
	"time"
)

func BenchmarkGCMConnector_BroadcastMessagesSingleWorker(b *testing.B) {
	throughput := throughputBroadcastMessages(b, 1)
	protocol.Info("Throughput for single worker: %8f msg/sec\n", throughput)
}

func BenchmarkGCMConnector_BroadcastMessagesMultipleWorkers(b *testing.B) {
	gomaxprocs := runtime.GOMAXPROCS(0)
	throughput := throughputBroadcastMessages(b, gomaxprocs)
	protocol.Info("Throughput for GOMAXPROCS (%v): %8f msg/sec\n", gomaxprocs, throughput)
}

func throughputBroadcastMessages(b *testing.B, nWorkers int) float64 {
	testutil.EnableDebugForMethod()
	protocol.Debug("b.N=%v\n", b.N)
	defer testutil.ResetDefaultRegistryHealthCheck()
	a := assert.New(b)

	dir, errTempDir := ioutil.TempDir("", "guble_benchmarking_gcm_test")
	defer func() {
		errRemove := os.RemoveAll(dir)
		if errRemove != nil {
			protocol.Err("Could not remove directory", errRemove)
		}
	}()
	a.NoError(errTempDir)

	args := Args{
		Listen:      "localhost:0",
		KVBackend:   "memory",
		MSBackend:   "file",
		StoragePath: dir,
		GcmEnable:   true,
		GcmApiKey:   "WILL BE OVERWRITTEN",
		GcmWorkers:  nWorkers}

	service := StartService(args)

	protocol.Debug("Overwriting the GCM Sender with a Mock")
	gcmConnector, ok := service.Modules()[2].(*gcm.GCMConnector)
	a.True(ok, "Modules[2] should be of type GCMConnector")
	gcmConnector.Sender = testutil.CreateGcmSender(
		testutil.CreateRoundTripperWithJsonResponse(http.StatusOK, testutil.CorrectGcmResponseMessageJSON, nil))

	service.Start()

	gomaxprocs := runtime.GOMAXPROCS(0)
	clients := make([]client.Client, 0, gomaxprocs)
	for clientID := 0; clientID < gomaxprocs; clientID++ {
		location := "ws://" + service.WebServer().GetAddr() + "/stream/user/" + strconv.Itoa(clientID)
		client, err := client.Open(location, "http://localhost/", 1000, true)
		a.NoError(err)
		clients = append(clients, client)
	}

	start := time.Now()
	b.ResetTimer()

	for clientID, c := range clients {
		go func(cID int, c client.Client) {
			for i := 0; i < b.N; i++ {
				protocol.Debug("send %d %d", cID, i)
				a.NoError(c.Send("/gcm/broadcast", "special offer", "{id:id}"))
			}
		}(clientID, c)
	}

	//TODO Cosmin: this delay should be eliminated after clean-shutdown is operational
	time.Sleep(2 * time.Second)

	// stop service (and wait for all the messages to be processed)
	service.StopGracePeriod = 10 * time.Second
	err := service.Stop()
	a.Nil(err)

	end := time.Now()
	b.StopTimer()

	return float64(b.N*gomaxprocs) / end.Sub(start).Seconds()
}
<commit_msg>disabling debug in benchmarks<commit_after>package gubled

import (
	"github.com/smancke/guble/client"
	"github.com/smancke/guble/gcm"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/testutil"

	"github.com/stretchr/testify/assert"

	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"testing"
	"time"
)

func BenchmarkGCMConnector_BroadcastMessagesSingleWorker(b *testing.B) {
	throughput := throughputBroadcastMessages(b, 1)
	protocol.Info("Throughput for single worker: %8f msg/sec\n", throughput)
}

func BenchmarkGCMConnector_BroadcastMessagesMultipleWorkers(b *testing.B) {
	gomaxprocs := runtime.GOMAXPROCS(0)
	throughput := throughputBroadcastMessages(b, gomaxprocs)
	protocol.Info("Throughput for GOMAXPROCS (%v): %8f msg/sec\n", gomaxprocs, throughput)
}

func throughputBroadcastMessages(b *testing.B, nWorkers int) float64 {
	//testutil.EnableDebugForMethod()
	protocol.Debug("b.N=%v\n", b.N)
	defer testutil.ResetDefaultRegistryHealthCheck()
	a := assert.New(b)

	dir, errTempDir := ioutil.TempDir("", "guble_benchmarking_gcm_test")
	defer func() {
		errRemove := os.RemoveAll(dir)
		if errRemove != nil {
			protocol.Err("Could not remove directory", errRemove)
		}
	}()
	a.NoError(errTempDir)

	args := Args{
		Listen:      "localhost:0",
		KVBackend:   "memory",
		MSBackend:   "file",
		StoragePath: dir,
		GcmEnable:   true,
		GcmApiKey:   "WILL BE OVERWRITTEN",
		GcmWorkers:  nWorkers}

	service := StartService(args)

	protocol.Debug("Overwriting the GCM Sender with a Mock")
	gcmConnector, ok := service.Modules()[2].(*gcm.GCMConnector)
	a.True(ok, "Modules[2] should be of type GCMConnector")
	gcmConnector.Sender = testutil.CreateGcmSender(
		testutil.CreateRoundTripperWithJsonResponse(http.StatusOK, testutil.CorrectGcmResponseMessageJSON, nil))

	service.Start()

	gomaxprocs := runtime.GOMAXPROCS(0)
	clients := make([]client.Client, 0, gomaxprocs)
	for clientID := 0; clientID < gomaxprocs; clientID++ {
		location := "ws://" + service.WebServer().GetAddr() + "/stream/user/" + strconv.Itoa(clientID)
		client, err := client.Open(location, "http://localhost/", 1000, true)
		a.NoError(err)
		clients = append(clients, client)
	}

	start := time.Now()
	b.ResetTimer()

	for clientID, c := range clients {
		go func(cID int, c client.Client) {
			for i := 0; i < b.N; i++ {
				protocol.Debug("send %d %d", cID, i)
				a.NoError(c.Send("/gcm/broadcast", "special offer", "{id:id}"))
			}
		}(clientID, c)
	}

	//TODO Cosmin: this delay should be eliminated after clean-shutdown is operational
	time.Sleep(2 * time.Second)

	// stop service (and wait for all the messages to be processed)
	service.StopGracePeriod = 10 * time.Second
	err := service.Stop()
	a.Nil(err)

	end := time.Now()
	b.StopTimer()

	return float64(b.N*gomaxprocs) / end.Sub(start).Seconds()
}
<|endoftext|>


<commit_before>package smd

import (
	"math"
	"os"
	"time"

	"github.com/ChristopherRabotin/gokalman"
	"github.com/ChristopherRabotin/ode"
	kitlog "github.com/go-kit/kit/log"
	"github.com/gonum/matrix/mat64"
)

const (
	transitionPhiOnly = false
)

// OrbitEstimate is an ode.Integrable which allows to propagate an orbit via its initial estimate.
type OrbitEstimate struct {
	Φ      *mat64.Dense  // STM
	Orbit  Orbit         // estimated orbit
	Perts  Perturbations // perturbations to account for
	StopDT time.Time     // end time of te integration
	dt     time.Time     // current time of the integration
	step   time.Duration // time step
	logger kitlog.Logger // logger
}

// GetState gets the state.
func (e *OrbitEstimate) GetState() []float64 {
	rΦ, cΦ := e.Φ.Dims()
	s := make([]float64, 6+rΦ*cΦ)
	R, V := e.Orbit.RV()
	s[0] = R[0]
	s[1] = R[1]
	s[2] = R[2]
	s[3] = V[0]
	s[4] = V[1]
	s[5] = V[2]
	// Add the components of Φ
	sIdx := 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			s[sIdx] = e.Φ.At(i, j)
			sIdx++
		}
	}
	return s
}

// SetState sets the next state at time t.
func (e *OrbitEstimate) SetState(t float64, s []float64) {
	R := []float64{s[0], s[1], s[2]}
	V := []float64{s[3], s[4], s[5]}
	e.Orbit = *NewOrbitFromRV(R, V, e.Orbit.Origin)
	// Extract the components of Φ
	sIdx := 6
	rΦ, cΦ := e.Φ.Dims()
	Φk20 := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φk20.Set(i, j, s[sIdx])
			sIdx++
		}
	}
	if transitionPhiOnly {
		// Compute the Φ for this transition
		var Φinv mat64.Dense
		if err := Φinv.Inverse(e.Φ); err != nil {
			panic("could not invert e.Φ")
		}
		e.Φ.Mul(Φk20, &Φinv)
	} else {
		e.Φ = Φk20
	}
	// Increment the time.
	e.dt = e.dt.Add(e.step)
}

// Stop returns whether we should stop the integration.
func (e *OrbitEstimate) Stop(t float64) bool {
	return e.dt.After(e.StopDT)
}

// State returns the latest state
func (e *OrbitEstimate) State() State {
	return State{e.dt, Spacecraft{}, e.Orbit, nil, nil}
}

// Func does the math. Returns a new state.
func (e *OrbitEstimate) Func(t float64, f []float64) (fDot []float64) {
	// XXX: Note that this function is very similar to Mission.Func for a Cartesian propagation.
	// *BUT* we need to add in all the components of Φ, since they have to be integrated too.
	rΦ, cΦ := e.Φ.Dims()
	fDot = make([]float64, 6+rΦ*cΦ) // init return vector
	// Re-create the orbit from the state.
	R := []float64{f[0], f[1], f[2]}
	V := []float64{f[3], f[4], f[5]}
	orbit := NewOrbitFromRV(R, V, e.Orbit.Origin)
	bodyAcc := -orbit.Origin.μ / math.Pow(orbit.RNorm(), 3)
	// d\vec{R}/dt
	fDot[0] = f[3]
	fDot[1] = f[4]
	fDot[2] = f[5]
	// d\vec{V}/dt
	fDot[3] = bodyAcc * f[0]
	fDot[4] = bodyAcc * f[1]
	fDot[5] = bodyAcc * f[2]

	pert := e.Perts.Perturb(*orbit, e.dt, Spacecraft{})
	for i := 0; i < 6; i++ {
		fDot[i] += pert[i]
	}

	// Extract the components of Φ
	fIdx := 6
	Φ := mat64.NewDense(rΦ, cΦ, nil)
	ΦDot := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φ.Set(i, j, f[fIdx])
			fIdx++
		}
	}

	// Compute the STM.
	A := mat64.NewDense(6, 6, nil)
	// Top right is Identity 3x3
	A.Set(0, 3, 1)
	A.Set(1, 4, 1)
	A.Set(2, 5, 1)
	// Bottom left is where the magix is.
	x := R[0]
	y := R[1]
	z := R[2]
	x2 := math.Pow(R[0], 2)
	y2 := math.Pow(R[1], 2)
	z2 := math.Pow(R[2], 2)
	r2 := x2 + y2 + z2
	r232 := math.Pow(r2, 3/2.)
	r252 := math.Pow(r2, 5/2.)
	// Add the body perturbations

	dAxDx := 3*e.Orbit.Origin.μ*x2/r252 - e.Orbit.Origin.μ/r232
	dAxDy := 3 * e.Orbit.Origin.μ * x * y / r252
	dAxDz := 3 * e.Orbit.Origin.μ * x * z / r252
	dAyDx := 3 * e.Orbit.Origin.μ * x * y / r252
	dAyDy := 3*e.Orbit.Origin.μ*y2/r252 - e.Orbit.Origin.μ/r232
	dAyDz := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDx := 3 * e.Orbit.Origin.μ * x * z / r252
	dAzDy := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDz := 3*e.Orbit.Origin.μ*z2/r252 - e.Orbit.Origin.μ/r232

	A.Set(3, 0, dAxDx)
	A.Set(4, 0, dAyDx)
	A.Set(5, 0, dAzDx)
	A.Set(3, 1, dAxDy)
	A.Set(4, 1, dAyDy)
	A.Set(5, 1, dAzDy)
	A.Set(3, 2, dAxDz)
	A.Set(4, 2, dAyDz)
	A.Set(5, 2, dAzDz)

	// Jn perturbations:
	if e.Perts.Jn > 1 {
		// Ai0 = \frac{\partial a}{\partial x}
		// Ai1 = \frac{\partial a}{\partial y}
		// Ai2 = \frac{\partial a}{\partial z}
		A30 := A.At(3, 0)
		A40 := A.At(4, 0)
		A50 := A.At(5, 0)
		A31 := A.At(3, 1)
		A41 := A.At(4, 1)
		A51 := A.At(5, 1)
		A32 := A.At(3, 2)
		A42 := A.At(4, 2)
		A52 := A.At(5, 2)

		// Notation simplification
		z3 := math.Pow(R[2], 3)
		z4 := math.Pow(R[2], 4)
		// Adding those fractions to avoid forgetting the trailing period which makes them floats.
		f32 := 3 / 2.
		f152 := 15 / 2.
		r272 := math.Pow(r2, 7/2.)
		r292 := math.Pow(r2, 9/2.)
		// J2
		j2fact := e.Orbit.Origin.J(2) * math.Pow(e.Orbit.Origin.Radius, 2) * e.Orbit.Origin.μ
		A30 += -f32 * j2fact * (35*x2*z2/r292 - 5*x2/r272 - 5*z2/r272 + 1/r252) //dAxDx
		A40 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAyDx
		A50 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)                      //dAzDx

		A31 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAxDy
		A41 += -f32 * j2fact * (35*y2*z2/r292 - 5*y2/r272 - 5*z2/r272 + 1/r252) // dAyDy
		A51 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)                      // dAzDy

		A32 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)        //dAxDz
		A42 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)        //dAyDz
		A52 += -f32 * j2fact * (35*z4/r292 - 30*z2/r272 + 3/r252) // dAzDz

		// J3
		if e.Perts.Jn > 2 {
			z5 := math.Pow(R[2], 5)
			r2112 := math.Pow(r2, 11/2.)
			f52 := 5 / 2.
			f1052 := 105 / 2.
			j3fact := e.Orbit.Origin.J(3) * math.Pow(e.Orbit.Origin.Radius, 3) * e.Orbit.Origin.μ
			A30 += -f52 * j3fact * (63*x2*z3/r2112 - 21*x2*z/r292 - 7*z3/r292 + 3*z/r272) //dAxDx
			A40 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAyDx
			A50 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272)               //dAzDx

			A31 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAxDy
			A41 += -f52 * j3fact * (63*y2*z3/r2112 - 21*y2*z/r292 - 7*z3/r292 + 3*z/r272) // dAyDy
			A51 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272)               // dAzDy

			A32 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272) //dAxDz
			A42 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272) //dAyDz
			A52 += -f52 * j3fact * (63*z5/r2112 - 70*z3/r292 + 15*z/r272)   // dAzDz
		}
		// \frac{\partial a}{\partial x}
		A.Set(3, 0, A30)
		A.Set(4, 0, A40)
		A.Set(5, 0, A50)
		// \partial a/\partial y
		A.Set(3, 1, A31)
		A.Set(4, 1, A41)
		A.Set(5, 1, A51)
		// \partial a/\partial z
		A.Set(3, 2, A32)
		A.Set(4, 2, A42)
		A.Set(5, 2, A52)
	}
	ΦDot.Mul(A, Φ)

	// Store ΦDot in fDot
	fIdx = 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			fDot[fIdx] = ΦDot.At(i, j)
			fIdx++
		}
	}
	return fDot
}

// PropagateUntil propagates until the given time is reached.
func (e *OrbitEstimate) PropagateUntil(dt time.Time) {
	e.StopDT = dt
	ode.NewRK4(0, e.step.Seconds(), e).Solve() // Blocking.
}

// NewOrbitEstimate returns a new Estimate of an orbit given the perturbations to be taken into account.
// The only supported state is [\vec{r} \vec{v}]T (for now at least).
func NewOrbitEstimate(n string, o Orbit, p Perturbations, epoch time.Time, step time.Duration) *OrbitEstimate {
	// The initial previous STM is identity.
	klog := kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stdout))
	klog = kitlog.With(klog, "estimate", n)
	stopDT := epoch
	// XXX: We add the step for consistency with Mission. Mission is broken: it skips the first step because the time addition
	// happens in the Stop function instead of the SetState function, the former being called at the start of the integration.
	return &OrbitEstimate{gokalman.DenseIdentity(6), o, p, stopDT, epoch.Add(step), step, klog}
}
<commit_msg>Remove gokalman dependency which I though I had previously removed<commit_after>package smd

import (
	"math"
	"os"
	"time"

	"github.com/ChristopherRabotin/ode"
	kitlog "github.com/go-kit/kit/log"
	"github.com/gonum/matrix/mat64"
)

const (
	transitionPhiOnly = false
)

// OrbitEstimate is an ode.Integrable which allows to propagate an orbit via its initial estimate.
type OrbitEstimate struct {
	Φ      *mat64.Dense  // STM
	Orbit  Orbit         // estimated orbit
	Perts  Perturbations // perturbations to account for
	StopDT time.Time     // end time of te integration
	dt     time.Time     // current time of the integration
	step   time.Duration // time step
	logger kitlog.Logger // logger
}

// GetState gets the state.
func (e *OrbitEstimate) GetState() []float64 {
	rΦ, cΦ := e.Φ.Dims()
	s := make([]float64, 6+rΦ*cΦ)
	R, V := e.Orbit.RV()
	s[0] = R[0]
	s[1] = R[1]
	s[2] = R[2]
	s[3] = V[0]
	s[4] = V[1]
	s[5] = V[2]
	// Add the components of Φ
	sIdx := 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			s[sIdx] = e.Φ.At(i, j)
			sIdx++
		}
	}
	return s
}

// SetState sets the next state at time t.
func (e *OrbitEstimate) SetState(t float64, s []float64) {
	R := []float64{s[0], s[1], s[2]}
	V := []float64{s[3], s[4], s[5]}
	e.Orbit = *NewOrbitFromRV(R, V, e.Orbit.Origin)
	// Extract the components of Φ
	sIdx := 6
	rΦ, cΦ := e.Φ.Dims()
	Φk20 := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φk20.Set(i, j, s[sIdx])
			sIdx++
		}
	}
	if transitionPhiOnly {
		// Compute the Φ for this transition
		var Φinv mat64.Dense
		if err := Φinv.Inverse(e.Φ); err != nil {
			panic("could not invert e.Φ")
		}
		e.Φ.Mul(Φk20, &Φinv)
	} else {
		e.Φ = Φk20
	}
	// Increment the time.
	e.dt = e.dt.Add(e.step)
}

// Stop returns whether we should stop the integration.
func (e *OrbitEstimate) Stop(t float64) bool {
	return e.dt.After(e.StopDT)
}

// State returns the latest state
func (e *OrbitEstimate) State() State {
	return State{e.dt, Spacecraft{}, e.Orbit, nil, nil}
}

// Func does the math. Returns a new state.
func (e *OrbitEstimate) Func(t float64, f []float64) (fDot []float64) {
	// XXX: Note that this function is very similar to Mission.Func for a Cartesian propagation.
	// *BUT* we need to add in all the components of Φ, since they have to be integrated too.
	rΦ, cΦ := e.Φ.Dims()
	fDot = make([]float64, 6+rΦ*cΦ) // init return vector
	// Re-create the orbit from the state.
	R := []float64{f[0], f[1], f[2]}
	V := []float64{f[3], f[4], f[5]}
	orbit := NewOrbitFromRV(R, V, e.Orbit.Origin)
	bodyAcc := -orbit.Origin.μ / math.Pow(orbit.RNorm(), 3)
	// d\vec{R}/dt
	fDot[0] = f[3]
	fDot[1] = f[4]
	fDot[2] = f[5]
	// d\vec{V}/dt
	fDot[3] = bodyAcc * f[0]
	fDot[4] = bodyAcc * f[1]
	fDot[5] = bodyAcc * f[2]

	pert := e.Perts.Perturb(*orbit, e.dt, Spacecraft{})
	for i := 0; i < 6; i++ {
		fDot[i] += pert[i]
	}

	// Extract the components of Φ
	fIdx := 6
	Φ := mat64.NewDense(rΦ, cΦ, nil)
	ΦDot := mat64.NewDense(rΦ, cΦ, nil)
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			Φ.Set(i, j, f[fIdx])
			fIdx++
		}
	}

	// Compute the STM.
	A := mat64.NewDense(6, 6, nil)
	// Top right is Identity 3x3
	A.Set(0, 3, 1)
	A.Set(1, 4, 1)
	A.Set(2, 5, 1)
	// Bottom left is where the magix is.
	x := R[0]
	y := R[1]
	z := R[2]
	x2 := math.Pow(R[0], 2)
	y2 := math.Pow(R[1], 2)
	z2 := math.Pow(R[2], 2)
	r2 := x2 + y2 + z2
	r232 := math.Pow(r2, 3/2.)
	r252 := math.Pow(r2, 5/2.)
	// Add the body perturbations

	dAxDx := 3*e.Orbit.Origin.μ*x2/r252 - e.Orbit.Origin.μ/r232
	dAxDy := 3 * e.Orbit.Origin.μ * x * y / r252
	dAxDz := 3 * e.Orbit.Origin.μ * x * z / r252
	dAyDx := 3 * e.Orbit.Origin.μ * x * y / r252
	dAyDy := 3*e.Orbit.Origin.μ*y2/r252 - e.Orbit.Origin.μ/r232
	dAyDz := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDx := 3 * e.Orbit.Origin.μ * x * z / r252
	dAzDy := 3 * e.Orbit.Origin.μ * y * z / r252
	dAzDz := 3*e.Orbit.Origin.μ*z2/r252 - e.Orbit.Origin.μ/r232

	A.Set(3, 0, dAxDx)
	A.Set(4, 0, dAyDx)
	A.Set(5, 0, dAzDx)
	A.Set(3, 1, dAxDy)
	A.Set(4, 1, dAyDy)
	A.Set(5, 1, dAzDy)
	A.Set(3, 2, dAxDz)
	A.Set(4, 2, dAyDz)
	A.Set(5, 2, dAzDz)

	// Jn perturbations:
	if e.Perts.Jn > 1 {
		// Ai0 = \frac{\partial a}{\partial x}
		// Ai1 = \frac{\partial a}{\partial y}
		// Ai2 = \frac{\partial a}{\partial z}
		A30 := A.At(3, 0)
		A40 := A.At(4, 0)
		A50 := A.At(5, 0)
		A31 := A.At(3, 1)
		A41 := A.At(4, 1)
		A51 := A.At(5, 1)
		A32 := A.At(3, 2)
		A42 := A.At(4, 2)
		A52 := A.At(5, 2)

		// Notation simplification
		z3 := math.Pow(R[2], 3)
		z4 := math.Pow(R[2], 4)
		// Adding those fractions to avoid forgetting the trailing period which makes them floats.
		f32 := 3 / 2.
		f152 := 15 / 2.
		r272 := math.Pow(r2, 7/2.)
		r292 := math.Pow(r2, 9/2.)
		// J2
		j2fact := e.Orbit.Origin.J(2) * math.Pow(e.Orbit.Origin.Radius, 2) * e.Orbit.Origin.μ
		A30 += -f32 * j2fact * (35*x2*z2/r292 - 5*x2/r272 - 5*z2/r272 + 1/r252) //dAxDx
		A40 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAyDx
		A50 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)                      //dAzDx

		A31 += -f152 * j2fact * (7*x*y*z2/r292 - x*y/r272)                      //dAxDy
		A41 += -f32 * j2fact * (35*y2*z2/r292 - 5*y2/r272 - 5*z2/r272 + 1/r252) // dAyDy
		A51 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)                      // dAzDy

		A32 += -f152 * j2fact * (7*x*z3/r292 - 3*x*z/r272)        //dAxDz
		A42 += -f152 * j2fact * (7*y*z3/r292 - 3*y*z/r272)        //dAyDz
		A52 += -f32 * j2fact * (35*z4/r292 - 30*z2/r272 + 3/r252) // dAzDz

		// J3
		if e.Perts.Jn > 2 {
			z5 := math.Pow(R[2], 5)
			r2112 := math.Pow(r2, 11/2.)
			f52 := 5 / 2.
			f1052 := 105 / 2.
			j3fact := e.Orbit.Origin.J(3) * math.Pow(e.Orbit.Origin.Radius, 3) * e.Orbit.Origin.μ
			A30 += -f52 * j3fact * (63*x2*z3/r2112 - 21*x2*z/r292 - 7*z3/r292 + 3*z/r272) //dAxDx
			A40 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAyDx
			A50 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272)               //dAzDx

			A31 += -f1052 * j3fact * (3*x*y*z3/r2112 - x*y*z/r292)                        //dAxDy
			A41 += -f52 * j3fact * (63*y2*z3/r2112 - 21*y2*z/r292 - 7*z3/r292 + 3*z/r272) // dAyDy
			A51 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272)               // dAzDy

			A32 += -f152 * j3fact * (21*x*z4/r2112 - 14*x*z2/r292 + x/r272) //dAxDz
			A42 += -f152 * j3fact * (21*y*z4/r2112 - 14*y*z2/r292 + y/r272) //dAyDz
			A52 += -f52 * j3fact * (63*z5/r2112 - 70*z3/r292 + 15*z/r272)   // dAzDz
		}
		// \frac{\partial a}{\partial x}
		A.Set(3, 0, A30)
		A.Set(4, 0, A40)
		A.Set(5, 0, A50)
		// \partial a/\partial y
		A.Set(3, 1, A31)
		A.Set(4, 1, A41)
		A.Set(5, 1, A51)
		// \partial a/\partial z
		A.Set(3, 2, A32)
		A.Set(4, 2, A42)
		A.Set(5, 2, A52)
	}
	ΦDot.Mul(A, Φ)

	// Store ΦDot in fDot
	fIdx = 6
	for i := 0; i < rΦ; i++ {
		for j := 0; j < cΦ; j++ {
			fDot[fIdx] = ΦDot.At(i, j)
			fIdx++
		}
	}
	return fDot
}

// PropagateUntil propagates until the given time is reached.
func (e *OrbitEstimate) PropagateUntil(dt time.Time) {
	e.StopDT = dt
	ode.NewRK4(0, e.step.Seconds(), e).Solve() // Blocking.
}

// NewOrbitEstimate returns a new Estimate of an orbit given the perturbations to be taken into account.
// The only supported state is [\vec{r} \vec{v}]T (for now at least).
func NewOrbitEstimate(n string, o Orbit, p Perturbations, epoch time.Time, step time.Duration) *OrbitEstimate {
	// The initial previous STM is identity.
	klog := kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stdout))
	klog = kitlog.With(klog, "estimate", n)
	stopDT := epoch
	// XXX: We add the step for consistency with Mission. Mission is broken: it skips the first step because the time addition
	// happens in the Stop function instead of the SetState function, the former being called at the start of the integration.
	return &OrbitEstimate{DenseIdentity(6), o, p, stopDT, epoch.Add(step), step, klog}
}
<|endoftext|>


<commit_before>// Package baseconv converts numeric strings between arbitrary radices
//
// The math is done manually, so values are not restricted by system maximums
package baseconv

import "fmt"

var ErrDuplicateCharacter = fmt.Errorf("duplicate character")

const chars64 = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

var (
	Base2,  _ = NewBaseMap(chars64[:2])
	Base8,  _ = NewBaseMap(chars64[:8])
	Base10, _ = NewBaseMap(chars64[:10])
	Base16, _ = NewBaseMap(chars64[:16])
	Base36, _ = NewBaseMap(chars64[:36])
	Base62, _ = NewBaseMap(chars64)
)

// BaseMap represents a pool of characters used for a certain radix
type BaseMap struct {
	radix  int
	values map[rune]int
	chars  map[int]rune
}

// Parse converts a numeric string to a Number
func (bm *BaseMap) Parse(input string) (*Number, error) {
	// Use len(input) even though it returns the number of bytes not runes.
	// Some memory might be wasted, but it's way faster than utf8.RuneCountInString (over 100x faster)
	values := make([]int, 0, len(input))

	// add the value of each digit to values
	for _, c := range input {
		value, ok := bm.values[c]

		// gotta check for that character validity ;)
		if !ok {
			return nil, fmt.Errorf("invalid character %c", c)
		}

		values = append(values, value)
	}

	return &Number{vals: values, bmap: bm}, nil
}

// Number represents a number with an arbitrary radix
type Number struct {
	vals []int
	bmap *BaseMap
}

// Format converts the number to a numeric string according to the given basemap
func (num *Number) Format(bm *BaseMap) string {
	radix, divisor := num.bmap.radix, bm.radix

	values := num.vals[:]
	places := len(values)
	str := make([]rune, 0, places)
	for places > 0 {
		// do manual division
		var remainder, place int
		for i := 0; i < places; i++ {
			dividend := remainder * radix + values[i]
			if dividend >= divisor {
				remainder = dividend % divisor
			} else if i == 0 {
				remainder = dividend
				continue
			}

			values[place] = dividend / divisor

			place++
		}

		places = place

		str = append(str, bm.chars[remainder])
	}

	// reverse the order of str
	l := len(str)
	for i := 0; i < l / 2; i++ {
		str[i], str[l-1-i] = str[l-1-i], str[i]
	}

	return string(str)
}

// NewBaseMap generates a base map
func NewBaseMap(chars string) (*BaseMap, error) {
	// use maps to make things faster
	vm := make(map[rune]int)
	cm := make(map[int]rune)

	var count int
	for i, c := range chars {
		_, ok := vm[c]
		if ok {
			return nil, ErrDuplicateCharacter
		}

		vm[c] = i
		cm[i] = c

		count++
	}

	return &BaseMap{radix: count, values: vm, chars: cm}, nil
}

// Convert converts the input from one radix to another
func Convert(input, fromBase, toBase string) (string, error) {
	bm1, err := NewBaseMap(fromBase)
	if err != nil {
		return "", err
	}

	bm2, err := NewBaseMap(toBase)
	if err != nil {
		return "", err
	}

	num, err := bm1.Parse(input)
	if err != nil {
		return "", err
	}

	return num.Format(bm2), nil
}
<commit_msg>add convenience function to return int64 value<commit_after>// Package baseconv converts numeric strings between arbitrary radices
//
// The math is done manually, so values are not restricted by system maximums
package baseconv

import (
	"fmt"
	"strconv"
)

var ErrDuplicateCharacter = fmt.Errorf("duplicate character")

const chars64 = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

var (
	Base2,  _ = NewBaseMap(chars64[:2])
	Base8,  _ = NewBaseMap(chars64[:8])
	Base10, _ = NewBaseMap(chars64[:10])
	Base16, _ = NewBaseMap(chars64[:16])
	Base36, _ = NewBaseMap(chars64[:36])
	Base62, _ = NewBaseMap(chars64)
)

// BaseMap represents a pool of characters used for a certain radix
type BaseMap struct {
	radix  int
	values map[rune]int
	chars  map[int]rune
}

// Parse converts a numeric string to a Number
func (bm *BaseMap) Parse(input string) (*Number, error) {
	// Use len(input) even though it returns the number of bytes not runes.
	// Some memory might be wasted, but it's way faster than utf8.RuneCountInString (over 100x faster)
	values := make([]int, 0, len(input))

	// add the value of each digit to values
	for _, c := range input {
		value, ok := bm.values[c]

		// gotta check for that character validity ;)
		if !ok {
			return nil, fmt.Errorf("invalid character %c", c)
		}

		values = append(values, value)
	}

	return &Number{vals: values, bmap: bm}, nil
}

// Number represents a number with an arbitrary radix
type Number struct {
	vals []int
	bmap *BaseMap
}

// Int64 returns the int64 value of the number as long as it is within 64 bits
func (num *Number) Int64() int64 {
	intVal, _ := strconv.ParseInt(num.Format(Base10), 10, 64)
	return intVal
}

// Format converts the number to a numeric string according to the given basemap
func (num *Number) Format(bm *BaseMap) string {
	radix, divisor := num.bmap.radix, bm.radix

	values := num.vals[:]
	places := len(values)
	str := make([]rune, 0, places)
	for places > 0 {
		// do manual division
		var remainder, place int
		for i := 0; i < places; i++ {
			dividend := remainder * radix + values[i]
			if dividend >= divisor {
				remainder = dividend % divisor
			} else if i == 0 {
				remainder = dividend
				continue
			}

			values[place] = dividend / divisor

			place++
		}

		places = place

		str = append(str, bm.chars[remainder])
	}

	// reverse the order of str
	l := len(str)
	for i := 0; i < l / 2; i++ {
		str[i], str[l-1-i] = str[l-1-i], str[i]
	}

	return string(str)
}

// NewBaseMap generates a base map
func NewBaseMap(chars string) (*BaseMap, error) {
	// use maps to make things faster
	vm := make(map[rune]int)
	cm := make(map[int]rune)

	var count int
	for i, c := range chars {
		_, ok := vm[c]
		if ok {
			return nil, ErrDuplicateCharacter
		}

		vm[c] = i
		cm[i] = c

		count++
	}

	return &BaseMap{radix: count, values: vm, chars: cm}, nil
}

// Convert converts the input from one radix to another
func Convert(input, fromBase, toBase string) (string, error) {
	bm1, err := NewBaseMap(fromBase)
	if err != nil {
		return "", err
	}

	bm2, err := NewBaseMap(toBase)
	if err != nil {
		return "", err
	}

	num, err := bm1.Parse(input)
	if err != nil {
		return "", err
	}

	return num.Format(bm2), nil
}
<|endoftext|>