<commit_before>// Copyright © 2017 Mikael Berthe <mikael@lilotux.net>
//
// Licensed under the MIT license.
// Please see the LICENSE file is this directory.

package cmd

import (
	"io/ioutil"
	"os"
	"strings"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	flag "github.com/spf13/pflag"

	"github.com/McKael/madon"
)

var statusPostFlags *flag.FlagSet

var statusOpts struct {
	statusID int64
	unset    bool

	// The following fields are used for the post/toot command
	visibility     string
	sensitive      bool
	spoiler        string
	inReplyToID    int64
	mediaIDs       string
	mediaFilePath  string
	textFilePath   string
	stdin          bool
	addMentions    bool
	sameVisibility bool

	// Used for several subcommands to limit the number of results
	limit, keep uint
	//sinceID, maxID int64
	all bool
}

func init() {
	RootCmd.AddCommand(statusCmd)

	// Subcommands
	statusCmd.AddCommand(statusSubcommands...)

	// Global flags
	statusCmd.PersistentFlags().Int64VarP(&statusOpts.statusID, "status-id", "s", 0, "Status ID number")
	statusCmd.PersistentFlags().UintVarP(&statusOpts.limit, "limit", "l", 0, "Limit number of API results")
	statusCmd.PersistentFlags().UintVarP(&statusOpts.keep, "keep", "k", 0, "Limit number of results")
	//statusCmd.PersistentFlags().Int64Var(&statusOpts.sinceID, "since-id", 0, "Request IDs greater than a value")
	//statusCmd.PersistentFlags().Int64Var(&statusOpts.maxID, "max-id", 0, "Request IDs less (or equal) than a value")
	statusCmd.PersistentFlags().BoolVar(&statusOpts.all, "all", false, "Fetch all results (for reblogged-by/favourited-by)")

	statusCmd.MarkPersistentFlagRequired("status-id")

	// Subcommand flags
	statusReblogSubcommand.Flags().BoolVar(&statusOpts.unset, "unset", false, "Unreblog the status")
	statusFavouriteSubcommand.Flags().BoolVar(&statusOpts.unset, "unset", false, "Remove the status from the favourites")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.sensitive, "sensitive", false, "Mark post as sensitive (NSFW)")
	statusPostSubcommand.Flags().StringVar(&statusOpts.visibility, "visibility", "", "Visibility (direct|private|unlisted|public)")
	statusPostSubcommand.Flags().StringVar(&statusOpts.spoiler, "spoiler", "", "Spoiler warning (CW)")
	statusPostSubcommand.Flags().StringVar(&statusOpts.mediaIDs, "media-ids", "", "Comma-separated list of media IDs")
	statusPostSubcommand.Flags().StringVarP(&statusOpts.mediaFilePath, "file", "f", "", "Media file name")
	statusPostSubcommand.Flags().StringVar(&statusOpts.textFilePath, "text-file", "", "Text file name (message content)")
	statusPostSubcommand.Flags().Int64VarP(&statusOpts.inReplyToID, "in-reply-to", "r", 0, "Status ID to reply to")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.stdin, "stdin", false, "Read message content from standard input")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.addMentions, "add-mentions", false, "Add mentions when replying")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.sameVisibility, "same-visibility", false, "Use same visibility as original message (for replies)")

	// Flag completion
	annotation := make(map[string][]string)
	annotation[cobra.BashCompCustom] = []string{"__madonctl_visibility"}

	statusPostSubcommand.Flags().Lookup("visibility").Annotations = annotation

	// This one will be used to check if the options were explicitly set or not
	statusPostFlags = statusPostSubcommand.Flags()
}

// statusCmd represents the status command
// This command does nothing without a subcommand
var statusCmd = &cobra.Command{
	Use:     "status --status-id ID subcommand",
	Aliases: []string{"st"},
	Short:   "Get status details",
	//Long:    `TBW...`, // TODO
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		// This is common to status and all status subcommands but "post"
		if statusOpts.statusID < 1 && cmd.Name() != "post" {
			return errors.New("missing status ID")
		}
		return madonInit(true)
	},
}

var statusSubcommands = []*cobra.Command{
	&cobra.Command{
		Use:     "show",
		Aliases: []string{"display"},
		Short:   "Get the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:   "context",
		Short: "Get the status context",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:   "card",
		Short: "Get the status card",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:   "reblogged-by",
		Short: "Display accounts which reblogged the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "favourited-by",
		Aliases: []string{"favorited-by"},
		Short:   "Display accounts which favourited the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "delete",
		Aliases: []string{"rm"},
		Short:   "Delete the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "mute-conversation",
		Aliases: []string{"mute"},
		Short:   "Mute the conversation containing the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "unmute-conversation",
		Aliases: []string{"unmute"},
		Short:   "Unmute the conversation containing the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	statusReblogSubcommand,
	statusFavouriteSubcommand,
	statusPostSubcommand,
}

var statusReblogSubcommand = &cobra.Command{
	Use:     "boost",
	Aliases: []string{"reblog"},
	Short:   "Boost (reblog) or unreblog the status",
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

var statusFavouriteSubcommand = &cobra.Command{
	Use:     "favourite",
	Aliases: []string{"favorite", "fave"},
	Short:   "Mark/unmark the status as favourite",
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

var statusPostSubcommand = &cobra.Command{
	Use:     "post",
	Aliases: []string{"toot", "pouet"},
	Short:   "Post a message (same as 'madonctl toot')",
	Example: `  madonctl status post --spoiler Warning "Hello, World"
  madonctl status toot --sensitive --file image.jpg Image
  madonctl status post --media-ids ID1,ID2,ID3 Image
  madonctl status toot --text-file message.txt
  madonctl status post --in-reply-to STATUSID "@user response"
  madonctl status post --in-reply-to STATUSID --add-mentions "response"
  echo "Hello from #madonctl" | madonctl status toot --stdin

The default visibility can be set in the configuration file with the option
'default_visibility' (or with an environmnent variable).`,
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

func statusSubcommandRunE(subcmd string, args []string) error {
	opt := statusOpts

	var obj interface{}
	var err error

	var limOpts *madon.LimitParams
	if opt.all || opt.limit > 0 /* || opt.sinceID > 0 || opt.maxID > 0 */ {
		limOpts = new(madon.LimitParams)
		limOpts.All = opt.all
	}

	if opt.limit > 0 {
		limOpts.Limit = int(opt.limit)
	}
	/*
		if opt.maxID > 0 {
			limOpts.MaxID = int64(opt.maxID)
		}
		if opt.sinceID > 0 {
			limOpts.SinceID = int64(opt.sinceID)
		}
	*/

	switch subcmd {
	case "show":
		var status *madon.Status
		status, err = gClient.GetStatus(opt.statusID)
		obj = status
	case "context":
		var context *madon.Context
		context, err = gClient.GetStatusContext(opt.statusID)
		obj = context
	case "card":
		var context *madon.Card
		context, err = gClient.GetStatusCard(opt.statusID)
		obj = context
	case "reblogged-by":
		var accountList []madon.Account
		accountList, err = gClient.GetStatusRebloggedBy(opt.statusID, limOpts)
		if opt.keep > 0 && len(accountList) > int(opt.keep) {
			accountList = accountList[:opt.keep]
		}
		obj = accountList
	case "favourited-by":
		var accountList []madon.Account
		accountList, err = gClient.GetStatusFavouritedBy(opt.statusID, limOpts)
		if opt.keep > 0 && len(accountList) > int(opt.keep) {
			accountList = accountList[:opt.keep]
		}
		obj = accountList
	case "delete":
		err = gClient.DeleteStatus(opt.statusID)
	case "boost":
		if opt.unset {
			err = gClient.UnreblogStatus(opt.statusID)
		} else {
			err = gClient.ReblogStatus(opt.statusID)
		}
	case "favourite":
		if opt.unset {
			err = gClient.UnfavouriteStatus(opt.statusID)
		} else {
			err = gClient.FavouriteStatus(opt.statusID)
		}
	case "mute-conversation":
		var s *madon.Status
		s, err = gClient.MuteConversation(opt.statusID)
		obj = s
	case "unmute-conversation":
		var s *madon.Status
		s, err = gClient.UnmuteConversation(opt.statusID)
		obj = s
	case "post": // toot
		var s *madon.Status
		text := strings.Join(args, " ")
		if opt.textFilePath != "" {
			var b []byte
			if b, err = ioutil.ReadFile(opt.textFilePath); err != nil {
				break
			}
			text = string(b)
		} else if opt.stdin {
			var b []byte
			if b, err = ioutil.ReadAll(os.Stdin); err != nil {
				break
			}
			text = string(b)
		}
		s, err = toot(text)
		obj = s
	default:
		return errors.New("statusSubcommand: internal error")
	}

	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}
	if obj == nil {
		return nil
	}

	p, err := getPrinter()
	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}
	return p.printObj(obj)
}
<commit_msg>Pin/Unpin support<commit_after>// Copyright © 2017 Mikael Berthe <mikael@lilotux.net>
//
// Licensed under the MIT license.
// Please see the LICENSE file is this directory.

package cmd

import (
	"io/ioutil"
	"os"
	"strings"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	flag "github.com/spf13/pflag"

	"github.com/McKael/madon"
)

var statusPostFlags *flag.FlagSet

var statusOpts struct {
	statusID int64
	unset    bool

	// The following fields are used for the post/toot command
	visibility     string
	sensitive      bool
	spoiler        string
	inReplyToID    int64
	mediaIDs       string
	mediaFilePath  string
	textFilePath   string
	stdin          bool
	addMentions    bool
	sameVisibility bool

	// Used for several subcommands to limit the number of results
	limit, keep uint
	//sinceID, maxID int64
	all bool
}

func init() {
	RootCmd.AddCommand(statusCmd)

	// Subcommands
	statusCmd.AddCommand(statusSubcommands...)

	// Global flags
	statusCmd.PersistentFlags().Int64VarP(&statusOpts.statusID, "status-id", "s", 0, "Status ID number")
	statusCmd.PersistentFlags().UintVarP(&statusOpts.limit, "limit", "l", 0, "Limit number of API results")
	statusCmd.PersistentFlags().UintVarP(&statusOpts.keep, "keep", "k", 0, "Limit number of results")
	//statusCmd.PersistentFlags().Int64Var(&statusOpts.sinceID, "since-id", 0, "Request IDs greater than a value")
	//statusCmd.PersistentFlags().Int64Var(&statusOpts.maxID, "max-id", 0, "Request IDs less (or equal) than a value")
	statusCmd.PersistentFlags().BoolVar(&statusOpts.all, "all", false, "Fetch all results (for reblogged-by/favourited-by)")

	statusCmd.MarkPersistentFlagRequired("status-id")

	// Subcommand flags
	statusReblogSubcommand.Flags().BoolVar(&statusOpts.unset, "unset", false, "Unreblog the status")
	statusFavouriteSubcommand.Flags().BoolVar(&statusOpts.unset, "unset", false, "Remove the status from the favourites")
	statusPinSubcommand.Flags().BoolVar(&statusOpts.unset, "unset", false, "Unpin the status")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.sensitive, "sensitive", false, "Mark post as sensitive (NSFW)")
	statusPostSubcommand.Flags().StringVar(&statusOpts.visibility, "visibility", "", "Visibility (direct|private|unlisted|public)")
	statusPostSubcommand.Flags().StringVar(&statusOpts.spoiler, "spoiler", "", "Spoiler warning (CW)")
	statusPostSubcommand.Flags().StringVar(&statusOpts.mediaIDs, "media-ids", "", "Comma-separated list of media IDs")
	statusPostSubcommand.Flags().StringVarP(&statusOpts.mediaFilePath, "file", "f", "", "Media file name")
	statusPostSubcommand.Flags().StringVar(&statusOpts.textFilePath, "text-file", "", "Text file name (message content)")
	statusPostSubcommand.Flags().Int64VarP(&statusOpts.inReplyToID, "in-reply-to", "r", 0, "Status ID to reply to")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.stdin, "stdin", false, "Read message content from standard input")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.addMentions, "add-mentions", false, "Add mentions when replying")
	statusPostSubcommand.Flags().BoolVar(&statusOpts.sameVisibility, "same-visibility", false, "Use same visibility as original message (for replies)")

	// Flag completion
	annotation := make(map[string][]string)
	annotation[cobra.BashCompCustom] = []string{"__madonctl_visibility"}

	statusPostSubcommand.Flags().Lookup("visibility").Annotations = annotation

	// This one will be used to check if the options were explicitly set or not
	statusPostFlags = statusPostSubcommand.Flags()
}

// statusCmd represents the status command
// This command does nothing without a subcommand
var statusCmd = &cobra.Command{
	Use:     "status --status-id ID subcommand",
	Aliases: []string{"st"},
	Short:   "Get status details",
	//Long:    `TBW...`, // TODO
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		// This is common to status and all status subcommands but "post"
		if statusOpts.statusID < 1 && cmd.Name() != "post" {
			return errors.New("missing status ID")
		}
		return madonInit(true)
	},
}

var statusSubcommands = []*cobra.Command{
	&cobra.Command{
		Use:     "show",
		Aliases: []string{"display"},
		Short:   "Get the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:   "context",
		Short: "Get the status context",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:   "card",
		Short: "Get the status card",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:   "reblogged-by",
		Short: "Display accounts which reblogged the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "favourited-by",
		Aliases: []string{"favorited-by"},
		Short:   "Display accounts which favourited the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "delete",
		Aliases: []string{"rm"},
		Short:   "Delete the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "mute-conversation",
		Aliases: []string{"mute"},
		Short:   "Mute the conversation containing the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	&cobra.Command{
		Use:     "unmute-conversation",
		Aliases: []string{"unmute"},
		Short:   "Unmute the conversation containing the status",
		RunE: func(cmd *cobra.Command, args []string) error {
			return statusSubcommandRunE(cmd.Name(), args)
		},
	},
	statusReblogSubcommand,
	statusFavouriteSubcommand,
	statusPinSubcommand,
	statusPostSubcommand,
}

var statusReblogSubcommand = &cobra.Command{
	Use:     "boost",
	Aliases: []string{"reblog"},
	Short:   "Boost (reblog) or unreblog the status",
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

var statusFavouriteSubcommand = &cobra.Command{
	Use:     "favourite",
	Aliases: []string{"favorite", "fave"},
	Short:   "Mark/unmark the status as favourite",
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

var statusPinSubcommand = &cobra.Command{
	Use:   "pin",
	Short: "Pin/unpin the status",
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

var statusPostSubcommand = &cobra.Command{
	Use:     "post",
	Aliases: []string{"toot", "pouet"},
	Short:   "Post a message (same as 'madonctl toot')",
	Example: `  madonctl status post --spoiler Warning "Hello, World"
  madonctl status toot --sensitive --file image.jpg Image
  madonctl status post --media-ids ID1,ID2,ID3 Image
  madonctl status toot --text-file message.txt
  madonctl status post --in-reply-to STATUSID "@user response"
  madonctl status post --in-reply-to STATUSID --add-mentions "response"
  echo "Hello from #madonctl" | madonctl status toot --stdin

The default visibility can be set in the configuration file with the option
'default_visibility' (or with an environmnent variable).`,
	RunE: func(cmd *cobra.Command, args []string) error {
		return statusSubcommandRunE(cmd.Name(), args)
	},
}

func statusSubcommandRunE(subcmd string, args []string) error {
	opt := statusOpts

	var obj interface{}
	var err error

	var limOpts *madon.LimitParams
	if opt.all || opt.limit > 0 /* || opt.sinceID > 0 || opt.maxID > 0 */ {
		limOpts = new(madon.LimitParams)
		limOpts.All = opt.all
	}

	if opt.limit > 0 {
		limOpts.Limit = int(opt.limit)
	}
	/*
		if opt.maxID > 0 {
			limOpts.MaxID = int64(opt.maxID)
		}
		if opt.sinceID > 0 {
			limOpts.SinceID = int64(opt.sinceID)
		}
	*/

	switch subcmd {
	case "show":
		var status *madon.Status
		status, err = gClient.GetStatus(opt.statusID)
		obj = status
	case "context":
		var context *madon.Context
		context, err = gClient.GetStatusContext(opt.statusID)
		obj = context
	case "card":
		var context *madon.Card
		context, err = gClient.GetStatusCard(opt.statusID)
		obj = context
	case "reblogged-by":
		var accountList []madon.Account
		accountList, err = gClient.GetStatusRebloggedBy(opt.statusID, limOpts)
		if opt.keep > 0 && len(accountList) > int(opt.keep) {
			accountList = accountList[:opt.keep]
		}
		obj = accountList
	case "favourited-by":
		var accountList []madon.Account
		accountList, err = gClient.GetStatusFavouritedBy(opt.statusID, limOpts)
		if opt.keep > 0 && len(accountList) > int(opt.keep) {
			accountList = accountList[:opt.keep]
		}
		obj = accountList
	case "delete":
		err = gClient.DeleteStatus(opt.statusID)
	case "boost":
		if opt.unset {
			err = gClient.UnreblogStatus(opt.statusID)
		} else {
			err = gClient.ReblogStatus(opt.statusID)
		}
	case "favourite":
		if opt.unset {
			err = gClient.UnfavouriteStatus(opt.statusID)
		} else {
			err = gClient.FavouriteStatus(opt.statusID)
		}
	case "pin":
		if opt.unset {
			err = gClient.UnpinStatus(opt.statusID)
		} else {
			err = gClient.PinStatus(opt.statusID)
		}
	case "mute-conversation":
		var s *madon.Status
		s, err = gClient.MuteConversation(opt.statusID)
		obj = s
	case "unmute-conversation":
		var s *madon.Status
		s, err = gClient.UnmuteConversation(opt.statusID)
		obj = s
	case "post": // toot
		var s *madon.Status
		text := strings.Join(args, " ")
		if opt.textFilePath != "" {
			var b []byte
			if b, err = ioutil.ReadFile(opt.textFilePath); err != nil {
				break
			}
			text = string(b)
		} else if opt.stdin {
			var b []byte
			if b, err = ioutil.ReadAll(os.Stdin); err != nil {
				break
			}
			text = string(b)
		}
		s, err = toot(text)
		obj = s
	default:
		return errors.New("statusSubcommand: internal error")
	}

	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}
	if obj == nil {
		return nil
	}

	p, err := getPrinter()
	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}
	return p.printObj(obj)
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"

	"github.com/exercism/cli/api"
	"github.com/exercism/cli/config"
	"github.com/exercism/cli/workspace"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

// submitCmd lets people upload a solution to the website.
var submitCmd = &cobra.Command{
	Use:     "submit FILE1 [FILE2 ...]",
	Aliases: []string{"s"},
	Short:   "Submit your solution to an exercise.",
	Long: `Submit your solution to an Exercism exercise.

    Call the command with the list of files you want to submit.
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		cfg := config.NewConfig()

		usrCfg := viper.New()
		usrCfg.AddConfigPath(cfg.Dir)
		usrCfg.SetConfigName("user")
		usrCfg.SetConfigType("json")
		// Ignore error. If the file doesn't exist, that is fine.
		_ = usrCfg.ReadInConfig()
		cfg.UserViperConfig = usrCfg

		v := viper.New()
		v.AddConfigPath(cfg.Dir)
		v.SetConfigName("cli")
		v.SetConfigType("json")
		// Ignore error. If the file doesn't exist, that is fine.
		_ = v.ReadInConfig()

		return runSubmit(cfg, cmd.Flags(), args)
	},
}

type submitContext struct {
	usrCfg *viper.Viper
	flags  *pflag.FlagSet
	args   []string
}

func runSubmit(cfg config.Config, flags *pflag.FlagSet, args []string) error {
	usrCfg := cfg.UserViperConfig

	if err := validateUserConfig(usrCfg); err != nil {
		return err
	}

	ctx := &submitContext{args: args, flags: flags, usrCfg: usrCfg}

	if err := ctx.sanitizeArgs(); err != nil {
		return err
	}

	exercise, err := ctx.exercise()
	if err != nil {
		return err
	}

	if err = ctx.migrateLegacyMetadata(exercise); err != nil {
		return err
	}

	metadata, err := ctx.metadata(exercise)
	if err != nil {
		return err
	}

	documents, err := ctx.documents(exercise)
	if err != nil {
		return err
	}

	if err := ctx.submitRequest(metadata, documents); err != nil {
		return err
	}

	ctx.printResult(metadata)
	return nil
}

func newSubmitContext(usrCfg *viper.Viper, args []string) (*submitContext, error) {
	if usrCfg.GetString("token") == "" {
		return nil, fmt.Errorf(
			msgWelcomePleaseConfigure,
			config.SettingsURL(usrCfg.GetString("apibaseurl")),
			BinaryName,
		)
	}
	if usrCfg.GetString("workspace") == "" {
		return nil, fmt.Errorf(msgRerunConfigure, BinaryName)
	}
	return &submitContext{args: args, usrCfg: usrCfg}, nil
}

func (s *submitContext) sanitizeArgs() error {
	for i, arg := range s.args {
		var err error
		arg, err = filepath.Abs(arg)
		if err != nil {
			return err
		}

		info, err := os.Lstat(arg)
		if err != nil {
			if os.IsNotExist(err) {
				msg := `

    The file you are trying to submit cannot be found.

        %s

        `
				return fmt.Errorf(msg, arg)
			}
			return err
		}
		if info.IsDir() {
			msg := `

    You are submitting a directory, which is not currently supported.

        %s

    Please change into the directory and provide the path to the file(s) you wish to submit

        %s submit FILENAME

            `
			return fmt.Errorf(msg, arg, BinaryName)
		}

		src, err := filepath.EvalSymlinks(arg)
		if err != nil {
			return err
		}
		s.args[i] = src
	}
	return nil
}

func (s *submitContext) exercise() (workspace.Exercise, error) {
	ws, err := workspace.New(s.usrCfg.GetString("workspace"))
	if err != nil {
		return workspace.Exercise{}, err
	}

	var exerciseDir string
	for _, arg := range s.args {
		dir, err := ws.ExerciseDir(arg)
		if err != nil {
			if workspace.IsMissingMetadata(err) {
				return workspace.Exercise{}, errors.New(msgMissingMetadata)
			}
			return workspace.Exercise{}, err
		}
		if exerciseDir != "" && dir != exerciseDir {
			msg := `

    You are submitting files belonging to different solutions.
    Please submit the files for one solution at a time.

        `
			return workspace.Exercise{}, errors.New(msg)
		}
		exerciseDir = dir
	}

	return workspace.NewExerciseFromDir(exerciseDir), nil
}

func (s *submitContext) migrateLegacyMetadata(exercise workspace.Exercise) error {
	migrationStatus, err := exercise.MigrateLegacyMetadataFile()
	if err != nil {
		return err
	}
	if verbose, _ := s.flags.GetBool("verbose"); verbose {
		fmt.Fprintf(Err, migrationStatus.String())
	}
	return nil
}

func (s *submitContext) metadata(exercise workspace.Exercise) (*workspace.ExerciseMetadata, error) {
	metadata, err := workspace.NewExerciseMetadata(exercise.Filepath())

	if err != nil {
		return nil, err
	}

	if exercise.Slug != metadata.Exercise {
		// TODO: error msg should suggest running future doctor command
		msg := `

	The exercise directory does not match exercise slug in metadata:

		expected '%[1]s' but got '%[2]s'

	Please rename the directory '%[1]s' to '%[2]s' and try again.

		`
		return nil, fmt.Errorf(msg, exercise.Slug, metadata.Exercise)
	}

	if !metadata.IsRequester {
		// TODO: add test
		msg := `

    The solution you are submitting is not connected to your account.
    Please re-download the exercise to make sure it has the data it needs.

        %s download --exercise=%s --track=%s

        `
		return nil, fmt.Errorf(msg, BinaryName, metadata.Exercise, metadata.Track)
	}
	return metadata, nil
}

func (s *submitContext) documents(exercise workspace.Exercise) ([]workspace.Document, error) {
	docs := make([]workspace.Document, 0, len(s.args))
	for _, file := range s.args {
		// Don't submit empty files
		info, err := os.Stat(file)
		if err != nil {
			return nil, err
		}
		const maxFileSize int64 = 65535
		if info.Size() >= maxFileSize {
			msg := `

      The submitted file '%s' is larger than the max allowed file size of %d bytes.
      Please reduce the size of the file and try again.

            `
			return nil, fmt.Errorf(msg, file, maxFileSize)
		}
		if info.Size() == 0 {

			msg := `

    WARNING: Skipping empty file
             %s

        `
			fmt.Fprintf(Err, msg, file)
			continue
		}
		doc, err := workspace.NewDocument(exercise.Filepath(), file)
		if err != nil {
			return nil, err
		}
		docs = append(docs, doc)
	}
	if len(docs) == 0 {
		msg := `

    No files found to submit.

        `
		return nil, errors.New(msg)
	}
	return docs, nil
}

func (s *submitContext) submitRequest(metadata *workspace.ExerciseMetadata, docs []workspace.Document) error {
	if metadata.ID == "" {
		return errors.New("id is empty")
	}
	if len(docs) == 0 {
		return errors.New("docs is empty")
	}

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	for _, doc := range docs {
		file, err := os.Open(doc.Filepath())
		if err != nil {
			return err
		}
		defer file.Close()

		part, err := writer.CreateFormFile("files[]", doc.Path())
		if err != nil {
			return err
		}
		_, err = io.Copy(part, file)
		if err != nil {
			return err
		}
	}
	if err := writer.Close(); err != nil {
		return err
	}

	client, err := api.NewClient(s.usrCfg.GetString("token"), s.usrCfg.GetString("apibaseurl"))
	if err != nil {
		return err
	}
	url := fmt.Sprintf("%s/solutions/%s", s.usrCfg.GetString("apibaseurl"), metadata.ID)
	req, err := client.NewRequest("PATCH", url, body)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", writer.FormDataContentType())

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusBadRequest {
		var jsonErrBody apiErrorMessage
		if err := json.NewDecoder(resp.Body).Decode(&jsonErrBody); err != nil {
			return fmt.Errorf("failed to parse error response - %s", err)
		}

		return fmt.Errorf(jsonErrBody.Error.Message)
	}

	bb := &bytes.Buffer{}
	_, err = bb.ReadFrom(resp.Body)
	if err != nil {
		return err
	}
	return nil
}

func (s *submitContext) printResult(metadata *workspace.ExerciseMetadata) {
	msg := `

    Your solution has been submitted successfully.
    %s
`
	suffix := "View it at:\n\n    "
	if metadata.AutoApprove && metadata.Team == "" {
		suffix = "You can complete the exercise and unlock the next core exercise at:\n"
	}
	fmt.Fprintf(Err, msg, suffix)
	fmt.Fprintf(Out, "    %s\n\n", metadata.URL)
}

func init() {
	RootCmd.AddCommand(submitCmd)
}

type apiErrorMessage struct {
	Error struct {
		Type    string `json:"type"`
		Message string `json:"message"`
	} `json:"error,omitempty"`
}
<commit_msg>Move sanitizeArgs to constructor<commit_after>package cmd

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"

	"github.com/exercism/cli/api"
	"github.com/exercism/cli/config"
	"github.com/exercism/cli/workspace"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

// submitCmd lets people upload a solution to the website.
var submitCmd = &cobra.Command{
	Use:     "submit FILE1 [FILE2 ...]",
	Aliases: []string{"s"},
	Short:   "Submit your solution to an exercise.",
	Long: `Submit your solution to an Exercism exercise.

    Call the command with the list of files you want to submit.
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		cfg := config.NewConfig()

		usrCfg := viper.New()
		usrCfg.AddConfigPath(cfg.Dir)
		usrCfg.SetConfigName("user")
		usrCfg.SetConfigType("json")
		// Ignore error. If the file doesn't exist, that is fine.
		_ = usrCfg.ReadInConfig()
		cfg.UserViperConfig = usrCfg

		v := viper.New()
		v.AddConfigPath(cfg.Dir)
		v.SetConfigName("cli")
		v.SetConfigType("json")
		// Ignore error. If the file doesn't exist, that is fine.
		_ = v.ReadInConfig()

		return runSubmit(cfg, cmd.Flags(), args)
	},
}

type submitContext struct {
	usrCfg *viper.Viper
	flags  *pflag.FlagSet
	args   []string
}

func runSubmit(cfg config.Config, flags *pflag.FlagSet, args []string) error {
	ctx, err := newSubmitContext(cfg.UserViperConfig, flags, args)
	if err != nil {
		return err
	}

	exercise, err := ctx.exercise()
	if err != nil {
		return err
	}

	if err = ctx.migrateLegacyMetadata(exercise); err != nil {
		return err
	}

	metadata, err := ctx.metadata(exercise)
	if err != nil {
		return err
	}

	documents, err := ctx.documents(exercise)
	if err != nil {
		return err
	}

	if err := ctx.submitRequest(metadata, documents); err != nil {
		return err
	}

	ctx.printResult(metadata)
	return nil
}

// newSubmitContext creates a submitContext and sanitizes the arguments.
func newSubmitContext(usrCfg *viper.Viper, flags *pflag.FlagSet, args []string) (*submitContext, error) {
	if usrCfg.GetString("token") == "" {
		return nil, fmt.Errorf(
			msgWelcomePleaseConfigure,
			config.SettingsURL(usrCfg.GetString("apibaseurl")),
			BinaryName,
		)
	}
	if usrCfg.GetString("workspace") == "" {
		return nil, fmt.Errorf(msgRerunConfigure, BinaryName)
	}

	ctx := &submitContext{usrCfg: usrCfg, flags: flags, args: args}
	return ctx, ctx.sanitizeArgs()
}

func (s *submitContext) sanitizeArgs() error {
	for i, arg := range s.args {
		var err error
		arg, err = filepath.Abs(arg)
		if err != nil {
			return err
		}

		info, err := os.Lstat(arg)
		if err != nil {
			if os.IsNotExist(err) {
				msg := `

    The file you are trying to submit cannot be found.

        %s

        `
				return fmt.Errorf(msg, arg)
			}
			return err
		}
		if info.IsDir() {
			msg := `

    You are submitting a directory, which is not currently supported.

        %s

    Please change into the directory and provide the path to the file(s) you wish to submit

        %s submit FILENAME

            `
			return fmt.Errorf(msg, arg, BinaryName)
		}

		src, err := filepath.EvalSymlinks(arg)
		if err != nil {
			return err
		}
		s.args[i] = src
	}
	return nil
}

func (s *submitContext) exercise() (workspace.Exercise, error) {
	ws, err := workspace.New(s.usrCfg.GetString("workspace"))
	if err != nil {
		return workspace.Exercise{}, err
	}

	var exerciseDir string
	for _, arg := range s.args {
		dir, err := ws.ExerciseDir(arg)
		if err != nil {
			if workspace.IsMissingMetadata(err) {
				return workspace.Exercise{}, errors.New(msgMissingMetadata)
			}
			return workspace.Exercise{}, err
		}
		if exerciseDir != "" && dir != exerciseDir {
			msg := `

    You are submitting files belonging to different solutions.
    Please submit the files for one solution at a time.

        `
			return workspace.Exercise{}, errors.New(msg)
		}
		exerciseDir = dir
	}

	return workspace.NewExerciseFromDir(exerciseDir), nil
}

func (s *submitContext) migrateLegacyMetadata(exercise workspace.Exercise) error {
	migrationStatus, err := exercise.MigrateLegacyMetadataFile()
	if err != nil {
		return err
	}
	if verbose, _ := s.flags.GetBool("verbose"); verbose {
		fmt.Fprintf(Err, migrationStatus.String())
	}
	return nil
}

func (s *submitContext) metadata(exercise workspace.Exercise) (*workspace.ExerciseMetadata, error) {
	metadata, err := workspace.NewExerciseMetadata(exercise.Filepath())

	if err != nil {
		return nil, err
	}

	if exercise.Slug != metadata.Exercise {
		// TODO: error msg should suggest running future doctor command
		msg := `

	The exercise directory does not match exercise slug in metadata:

		expected '%[1]s' but got '%[2]s'

	Please rename the directory '%[1]s' to '%[2]s' and try again.

		`
		return nil, fmt.Errorf(msg, exercise.Slug, metadata.Exercise)
	}

	if !metadata.IsRequester {
		// TODO: add test
		msg := `

    The solution you are submitting is not connected to your account.
    Please re-download the exercise to make sure it has the data it needs.

        %s download --exercise=%s --track=%s

        `
		return nil, fmt.Errorf(msg, BinaryName, metadata.Exercise, metadata.Track)
	}
	return metadata, nil
}

func (s *submitContext) documents(exercise workspace.Exercise) ([]workspace.Document, error) {
	docs := make([]workspace.Document, 0, len(s.args))
	for _, file := range s.args {
		// Don't submit empty files
		info, err := os.Stat(file)
		if err != nil {
			return nil, err
		}
		const maxFileSize int64 = 65535
		if info.Size() >= maxFileSize {
			msg := `

      The submitted file '%s' is larger than the max allowed file size of %d bytes.
      Please reduce the size of the file and try again.

            `
			return nil, fmt.Errorf(msg, file, maxFileSize)
		}
		if info.Size() == 0 {

			msg := `

    WARNING: Skipping empty file
             %s

        `
			fmt.Fprintf(Err, msg, file)
			continue
		}
		doc, err := workspace.NewDocument(exercise.Filepath(), file)
		if err != nil {
			return nil, err
		}
		docs = append(docs, doc)
	}
	if len(docs) == 0 {
		msg := `

    No files found to submit.

        `
		return nil, errors.New(msg)
	}
	return docs, nil
}

func (s *submitContext) submitRequest(metadata *workspace.ExerciseMetadata, docs []workspace.Document) error {
	if metadata.ID == "" {
		return errors.New("id is empty")
	}
	if len(docs) == 0 {
		return errors.New("docs is empty")
	}

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	for _, doc := range docs {
		file, err := os.Open(doc.Filepath())
		if err != nil {
			return err
		}
		defer file.Close()

		part, err := writer.CreateFormFile("files[]", doc.Path())
		if err != nil {
			return err
		}
		_, err = io.Copy(part, file)
		if err != nil {
			return err
		}
	}
	if err := writer.Close(); err != nil {
		return err
	}

	client, err := api.NewClient(s.usrCfg.GetString("token"), s.usrCfg.GetString("apibaseurl"))
	if err != nil {
		return err
	}
	url := fmt.Sprintf("%s/solutions/%s", s.usrCfg.GetString("apibaseurl"), metadata.ID)
	req, err := client.NewRequest("PATCH", url, body)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", writer.FormDataContentType())

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusBadRequest {
		var jsonErrBody apiErrorMessage
		if err := json.NewDecoder(resp.Body).Decode(&jsonErrBody); err != nil {
			return fmt.Errorf("failed to parse error response - %s", err)
		}

		return fmt.Errorf(jsonErrBody.Error.Message)
	}

	bb := &bytes.Buffer{}
	_, err = bb.ReadFrom(resp.Body)
	if err != nil {
		return err
	}
	return nil
}

func (s *submitContext) printResult(metadata *workspace.ExerciseMetadata) {
	msg := `

    Your solution has been submitted successfully.
    %s
`
	suffix := "View it at:\n\n    "
	if metadata.AutoApprove && metadata.Team == "" {
		suffix = "You can complete the exercise and unlock the next core exercise at:\n"
	}
	fmt.Fprintf(Err, msg, suffix)
	fmt.Fprintf(Out, "    %s\n\n", metadata.URL)
}

func init() {
	RootCmd.AddCommand(submitCmd)
}

type apiErrorMessage struct {
	Error struct {
		Type    string `json:"type"`
		Message string `json:"message"`
	} `json:"error,omitempty"`
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
	"github.com/vbauerster/mpb"
	"golang.org/x/sync/errgroup"

	"github.com/Shopify/themekit/cmd/ystore"
	"github.com/Shopify/themekit/kit"
)

const settingsDataKey = "config/settings_data.json"

var uploadCmd = &cobra.Command{
	Use:   "upload <filenames>",
	Short: "Upload theme file(s) to shopify",
	Long: `Upload will upload specific files to shopify servers if provided file names.
If no filenames are provided then upload will upload every file in the project
to shopify.

For more documentation please see http://shopify.github.io/themekit/commands/#upload
`,
	PreRunE:  arbiter.generateThemeClients,
	RunE:     arbiter.forEachClient(deploy(false)),
	PostRunE: arbiter.forEachClient(uploadSettingsData),
}

func deploy(destructive bool) arbitratedCmd {
	return func(client kit.ThemeClient, filenames []string) error {
		if client.Config.ReadOnly {
			return fmt.Errorf("[%s] environment is reaonly", green(client.Config.Environment))
		}

		actions, err := arbiter.generateAssetActions(client, filenames, destructive)
		if err != nil {
			return err
		}

		if err := arbiter.preflightCheck(actions, destructive); err != nil {
			return err
		}

		var deployGroup errgroup.Group
		bar := arbiter.newProgressBar(len(actions), client.Config.Environment)
		for key, action := range actions {
			shouldPerform := arbiter.force || arbiter.manifest.Should(action.event, action.asset, client.Config.Environment)
			// pretend we did the settings data and we will do it last
			if !shouldPerform || key == settingsDataKey {
				if arbiter.verbose {
					stdOut.Printf(
						"[%s] skipping %s of %s",
						green(client.Config.Environment),
						yellow(action.event),
						blue(action.asset.Key),
					)
				}

				incBar(bar)
				continue
			}
			action := action
			deployGroup.Go(func() error {
				if err := perform(client, action.asset, action.event, bar); err != nil {
					stdErr.Printf("[%s] %s", green(client.Config.Environment), err)
				}
				return nil
			})
		}
		return deployGroup.Wait()
	}
}

func incBar(bar *mpb.Bar) {
	if bar != nil {
		defer bar.Incr(1)
	}
}

func perform(client kit.ThemeClient, asset kit.Asset, event kit.EventType, bar *mpb.Bar) error {
	defer incBar(bar)

	var (
		resp    *kit.ShopifyResponse
		err     error
		version string
	)

	if arbiter.force {
		resp, err = client.Perform(asset, event)
	} else if version, _, err = arbiter.manifest.Get(asset.Key, client.Config.Environment); err == nil {
		resp, err = client.PerformStrict(asset, event, version)
	}

	if err != nil {
		return err
	} else if arbiter.verbose {
		stdOut.Printf(
			"[%s] Successfully performed %s on file %s from %s",
			green(client.Config.Environment),
			green(resp.EventType),
			blue(resp.Asset.Key),
			yellow(resp.Host),
		)
	}

	if event == kit.Remove {
		if err := arbiter.manifest.Delete(resp.Asset.Key, client.Config.Environment); err != nil && err != ystore.ErrorCollectionNotFound {
			return err
		}
	} else {
		checksum, _ := asset.CheckSum()
		if err := arbiter.manifest.Set(resp.Asset.Key, client.Config.Environment, resp.Asset.UpdatedAt, checksum); err != nil {
			return err
		}
	}

	return nil
}

func uploadSettingsData(client kit.ThemeClient, files []string) error {
	if len(files) == 0 {
		if actions, err := arbiter.generateAssetActions(client, files, false); err != nil {
			return err
		} else if _, found := actions[settingsDataKey]; !found {
			return nil
		}
	} else if i := indexOf(len(files), func(i int) bool { return files[i] == settingsDataKey }); i == -1 {
		return nil
	}

	asset, err := client.LocalAsset(settingsDataKey)
	if err != nil {
		return err
	}
	return perform(client, asset, kit.Update, nil)
}

func indexOf(count int, cb func(i int) bool) int {
	for i := 0; i < count; i++ {
		if cb(i) {
			return i
		}
	}
	return -1
}
<commit_msg>Typo fix at cmd/upload.go<commit_after>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
	"github.com/vbauerster/mpb"
	"golang.org/x/sync/errgroup"

	"github.com/Shopify/themekit/cmd/ystore"
	"github.com/Shopify/themekit/kit"
)

const settingsDataKey = "config/settings_data.json"

var uploadCmd = &cobra.Command{
	Use:   "upload <filenames>",
	Short: "Upload theme file(s) to shopify",
	Long: `Upload will upload specific files to shopify servers if provided file names.
If no filenames are provided then upload will upload every file in the project
to shopify.

For more documentation please see http://shopify.github.io/themekit/commands/#upload
`,
	PreRunE:  arbiter.generateThemeClients,
	RunE:     arbiter.forEachClient(deploy(false)),
	PostRunE: arbiter.forEachClient(uploadSettingsData),
}

func deploy(destructive bool) arbitratedCmd {
	return func(client kit.ThemeClient, filenames []string) error {
		if client.Config.ReadOnly {
			return fmt.Errorf("[%s] environment is readonly", green(client.Config.Environment))
		}

		actions, err := arbiter.generateAssetActions(client, filenames, destructive)
		if err != nil {
			return err
		}

		if err := arbiter.preflightCheck(actions, destructive); err != nil {
			return err
		}

		var deployGroup errgroup.Group
		bar := arbiter.newProgressBar(len(actions), client.Config.Environment)
		for key, action := range actions {
			shouldPerform := arbiter.force || arbiter.manifest.Should(action.event, action.asset, client.Config.Environment)
			// pretend we did the settings data and we will do it last
			if !shouldPerform || key == settingsDataKey {
				if arbiter.verbose {
					stdOut.Printf(
						"[%s] skipping %s of %s",
						green(client.Config.Environment),
						yellow(action.event),
						blue(action.asset.Key),
					)
				}

				incBar(bar)
				continue
			}
			action := action
			deployGroup.Go(func() error {
				if err := perform(client, action.asset, action.event, bar); err != nil {
					stdErr.Printf("[%s] %s", green(client.Config.Environment), err)
				}
				return nil
			})
		}
		return deployGroup.Wait()
	}
}

func incBar(bar *mpb.Bar) {
	if bar != nil {
		defer bar.Incr(1)
	}
}

func perform(client kit.ThemeClient, asset kit.Asset, event kit.EventType, bar *mpb.Bar) error {
	defer incBar(bar)

	var (
		resp    *kit.ShopifyResponse
		err     error
		version string
	)

	if arbiter.force {
		resp, err = client.Perform(asset, event)
	} else if version, _, err = arbiter.manifest.Get(asset.Key, client.Config.Environment); err == nil {
		resp, err = client.PerformStrict(asset, event, version)
	}

	if err != nil {
		return err
	} else if arbiter.verbose {
		stdOut.Printf(
			"[%s] Successfully performed %s on file %s from %s",
			green(client.Config.Environment),
			green(resp.EventType),
			blue(resp.Asset.Key),
			yellow(resp.Host),
		)
	}

	if event == kit.Remove {
		if err := arbiter.manifest.Delete(resp.Asset.Key, client.Config.Environment); err != nil && err != ystore.ErrorCollectionNotFound {
			return err
		}
	} else {
		checksum, _ := asset.CheckSum()
		if err := arbiter.manifest.Set(resp.Asset.Key, client.Config.Environment, resp.Asset.UpdatedAt, checksum); err != nil {
			return err
		}
	}

	return nil
}

func uploadSettingsData(client kit.ThemeClient, files []string) error {
	if len(files) == 0 {
		if actions, err := arbiter.generateAssetActions(client, files, false); err != nil {
			return err
		} else if _, found := actions[settingsDataKey]; !found {
			return nil
		}
	} else if i := indexOf(len(files), func(i int) bool { return files[i] == settingsDataKey }); i == -1 {
		return nil
	}

	asset, err := client.LocalAsset(settingsDataKey)
	if err != nil {
		return err
	}
	return perform(client, asset, kit.Update, nil)
}

func indexOf(count int, cb func(i int) bool) int {
	for i := 0; i < count; i++ {
		if cb(i) {
			return i
		}
	}
	return -1
}
<|endoftext|>


<commit_before>package models

import (
	"database/sql"
	"errors"
	"fmt"
	sq "github.com/Masterminds/squirrel"
	"github.com/guregu/null"
	"net/http"
	"strings"
)

// This is the "JSON" struct that appears in the array returned by getRecentSeeks
type SeeksItem struct {
	KeyID                int         `json:"keyId"`
	CreationDate         null.Time   `json:"creationDate"`
	LastModificationDate null.Time   `json:"lastModificationDate"`
	Title                string      `json:"title"`
	Description          null.String `json:"description"` // expect to be truncated
	UserID               int         `json:"userId"`
	Username             null.String `json:"username"`
	SavedSearchID        null.Int    `json:"savedSearchId"`
	NotifyEnabled        null.Bool   `json:"notifyEnabled"`
	Status               null.String `json:"status"`
}

// TODO Maybe use the same struct for both of these? The only difference is truncation of the descrip
// Returned by a function returning only one seek (usually by ID)
type Seek struct {
	KeyID                int         `json:"keyId"`
	CreationDate         null.Time   `json:"creationDate"`
	LastModificationDate null.Time   `json:"lastModificationDate"`
	Title                string      `json:"title"`
	Description          null.String `json:"description"`
	UserID               int         `json:"userId"`
	Username             null.String `json:"username"`
	SavedSearchID        null.Int    `json:"savedSearchId"`
	NotifyEnabled        null.Bool   `json:"notifyEnabled"`
	Status               null.String `json:"status"`
}

type seekQuery struct {
	Query            string
	OnlyMine         bool
	TruncationLength int
	Limit            uint64
	UserID           int
}

func NewSeekQuery() *seekQuery {
	q := new(seekQuery)
	q.TruncationLength = defaultTruncationLength
	q.Limit = defaultNumResults
	return q
}

// Returns the most recent count seeks, based on original date created.
// If queryStr is nonempty, filters that every returned item must have every word in either title or description
// On error, returns an error and the HTTP code associated with that error.
func ReadSeeks(db *sql.DB, query *seekQuery) ([]*SeeksItem, error, int) {
	// Create seeks statement
	stmt := psql.
		Select("seeks.key_id", "seeks.creation_date", "seeks.last_modification_date",
			"title", fmt.Sprintf("left(description, %d)", query.TruncationLength),
			"user_id", "users.net_id", "saved_search_id", "notify_enabled", "status").
		From("seeks").
		Where("seeks.is_active=true").
		LeftJoin("users ON seeks.user_id = users.key_id")

	for _, word := range strings.Fields(query.Query) {
		stmt = stmt.Where("(lower(seeks.title) LIKE lower(?) OR lower(seeks.description) LIKE lower(?))", fmt.Sprint("%", word, "%"), fmt.Sprint("%", word, "%"))
	}

	if query.UserID == 0 && query.OnlyMine {
		return nil, errors.New("Unauthenticated user attempted to view profile data"), http.StatusUnauthorized
	}

	if query.OnlyMine {
		stmt = stmt.Where(sq.Eq{"user_id": query.UserID})
	}

	stmt = stmt.OrderBy("seeks.creation_date DESC")

	if query.Limit <= maxNumResults {
		stmt = stmt.Limit(query.Limit)
	} else {
		stmt = stmt.Limit(maxNumResults)
	}

	// Query db
	rows, err := stmt.RunWith(db).Query()
	if err != nil {
		return nil, err, http.StatusInternalServerError
	}
	defer rows.Close()

	// Populate seek structs
	seeks := make([]*SeeksItem, 0)
	for rows.Next() {
		s := new(SeeksItem)
		err := rows.Scan(&s.KeyID, &s.CreationDate, &s.LastModificationDate,
			&s.Title, &s.Description, &s.UserID, &s.Username, &s.SavedSearchID,
			&s.NotifyEnabled, &s.Status)
		if err != nil {
			return nil, err, http.StatusInternalServerError
		}
		seeks = append(seeks, s)
	}
	if err = rows.Err(); err != nil {
		return nil, err, http.StatusInternalServerError
	}

	return seeks, nil, http.StatusOK
}

// Returns the most recent count seeks, based on original date created. On error
// returns an error and the HTTP code associated with that error.
func ReadSeek(db *sql.DB, id string) (Seek, error, int) {
	var seek Seek

	// Create seek query
	query := psql.
		Select("seeks.key_id", "seeks.creation_date",
			"seeks.last_modification_date", "title", "description", "user_id",
			"users.net_id", "saved_search_id", "notify_enabled", "status").
		From("seeks").
		Where("seeks.is_active=true").
		LeftJoin("users ON seeks.user_id = users.key_id").
		Where(sq.Eq{"seeks.key_id": id})

	// Query db for seek
	rows, err := query.RunWith(db).Query()
	if err != nil {
		return seek, err, http.StatusInternalServerError
	}
	defer rows.Close()

	// Populate seek struct
	rows.Next()
	err = rows.Scan(&seek.KeyID, &seek.CreationDate, &seek.LastModificationDate,
		&seek.Title, &seek.Description, &seek.UserID, &seek.Username, &seek.SavedSearchID,
		&seek.NotifyEnabled, &seek.Status)
	if err == sql.ErrNoRows {
		return seek, err, http.StatusNotFound
	} else if err != nil {
		return seek, err, http.StatusInternalServerError
	}

	return seek, nil, http.StatusOK
}

// Inserts the given seek (belonging to userId) into the database. Returns
// seek with its new KeyID added.
func CreateSeek(db *sql.DB, seek Seek, userId int) (Seek, error, int) {
	seek.UserID = userId

	// Insert seek
	stmt := psql.Insert("seeks").
		Columns("title", "description", "user_id", "saved_search_id",
			"notify_enabled", "status").
		Values(seek.Title, seek.Description, userId, seek.SavedSearchID,
			seek.NotifyEnabled, seek.Status).
		Suffix("RETURNING key_id, creation_date")

	// Query db for seek
	rows, err := stmt.RunWith(db).Query()
	if err != nil {
		return seek, err, http.StatusInternalServerError
	}
	defer rows.Close()

	// Populate seek struct
	rows.Next()
	err = rows.Scan(&seek.KeyID, &seek.CreationDate)
	if err != nil {
		return seek, err, http.StatusInternalServerError
	}

	return seek, nil, http.StatusCreated
}

// Overwrites the seek in the database with the given id with the given seek
// (belonging to userId). Returns the updated seek.
func UpdateSeek(db *sql.DB, id string, seek Seek, userId int) (error, int) {
	seek.UserID = userId

	// Update seek
	stmt := psql.Update("seeks").
		SetMap(map[string]interface{}{
			"title":           seek.Title,
			"description":     seek.Description,
			"saved_search_id": seek.SavedSearchID,
			"notify_enabled":  seek.NotifyEnabled}).
		Where(sq.Eq{"seeks.key_id": id,
			"seeks.user_id": userId})

	// Query db for seek
	result, err := stmt.RunWith(db).Exec()
	return getUpdateResultCode(result, err)
}

// Deletes the seek in the database with the given id with the given seek
// (belonging to userId).
func DeleteSeek(db *sql.DB, id string, userId int) (error, int) {
	// Update seek
	stmt := psql.Delete("seeks").
		Where(sq.Eq{"seeks.key_id": id,
			"seeks.user_id": userId})

	// Query db for seek
	result, err := stmt.RunWith(db).Exec()
	return getUpdateResultCode(result, err)
}
<commit_msg>Remove unnecessary struct<commit_after>package models

import (
	"database/sql"
	"errors"
	"fmt"
	sq "github.com/Masterminds/squirrel"
	"github.com/guregu/null"
	"net/http"
	"strings"
)

// Returned by a function returning only one seek (usually by ID)
type Seek struct {
	KeyID                int         `json:"keyId"`
	CreationDate         null.Time   `json:"creationDate"`
	LastModificationDate null.Time   `json:"lastModificationDate"`
	Title                string      `json:"title"`
	Description          null.String `json:"description"`
	UserID               int         `json:"userId"`
	Username             null.String `json:"username"`
	SavedSearchID        null.Int    `json:"savedSearchId"`
	NotifyEnabled        null.Bool   `json:"notifyEnabled"`
	Status               null.String `json:"status"`
}

type seekQuery struct {
	Query            string
	OnlyMine         bool
	TruncationLength int
	Limit            uint64
	UserID           int
}

func NewSeekQuery() *seekQuery {
	q := new(seekQuery)
	q.TruncationLength = defaultTruncationLength
	q.Limit = defaultNumResults
	return q
}

// Returns the most recent count seeks, based on original date created.
// If queryStr is nonempty, filters that every returned item must have every word in either title or description
// On error, returns an error and the HTTP code associated with that error.
func ReadSeeks(db *sql.DB, query *seekQuery) ([]*Seek, error, int) {
	// Create seeks statement
	stmt := psql.
		Select("seeks.key_id", "seeks.creation_date", "seeks.last_modification_date",
			"title", fmt.Sprintf("left(description, %d)", query.TruncationLength),
			"user_id", "users.net_id", "saved_search_id", "notify_enabled", "status").
		From("seeks").
		Where("seeks.is_active=true").
		LeftJoin("users ON seeks.user_id = users.key_id")

	for _, word := range strings.Fields(query.Query) {
		stmt = stmt.Where("(lower(seeks.title) LIKE lower(?) OR lower(seeks.description) LIKE lower(?))", fmt.Sprint("%", word, "%"), fmt.Sprint("%", word, "%"))
	}

	if query.UserID == 0 && query.OnlyMine {
		return nil, errors.New("Unauthenticated user attempted to view profile data"), http.StatusUnauthorized
	}

	if query.OnlyMine {
		stmt = stmt.Where(sq.Eq{"user_id": query.UserID})
	}

	stmt = stmt.OrderBy("seeks.creation_date DESC")

	if query.Limit <= maxNumResults {
		stmt = stmt.Limit(query.Limit)
	} else {
		stmt = stmt.Limit(maxNumResults)
	}

	// Query db
	rows, err := stmt.RunWith(db).Query()
	if err != nil {
		return nil, err, http.StatusInternalServerError
	}
	defer rows.Close()

	// Populate seek structs
	seeks := make([]*Seek, 0)
	for rows.Next() {
		s := new(Seek)
		err := rows.Scan(&s.KeyID, &s.CreationDate, &s.LastModificationDate,
			&s.Title, &s.Description, &s.UserID, &s.Username, &s.SavedSearchID,
			&s.NotifyEnabled, &s.Status)
		if err != nil {
			return nil, err, http.StatusInternalServerError
		}
		seeks = append(seeks, s)
	}
	if err = rows.Err(); err != nil {
		return nil, err, http.StatusInternalServerError
	}

	return seeks, nil, http.StatusOK
}

// Returns the most recent count seeks, based on original date created. On error
// returns an error and the HTTP code associated with that error.
func ReadSeek(db *sql.DB, id string) (Seek, error, int) {
	var seek Seek

	// Create seek query
	query := psql.
		Select("seeks.key_id", "seeks.creation_date",
			"seeks.last_modification_date", "title", "description", "user_id",
			"users.net_id", "saved_search_id", "notify_enabled", "status").
		From("seeks").
		Where("seeks.is_active=true").
		LeftJoin("users ON seeks.user_id = users.key_id").
		Where(sq.Eq{"seeks.key_id": id})

	// Query db for seek
	rows, err := query.RunWith(db).Query()
	if err != nil {
		return seek, err, http.StatusInternalServerError
	}
	defer rows.Close()

	// Populate seek struct
	rows.Next()
	err = rows.Scan(&seek.KeyID, &seek.CreationDate, &seek.LastModificationDate,
		&seek.Title, &seek.Description, &seek.UserID, &seek.Username, &seek.SavedSearchID,
		&seek.NotifyEnabled, &seek.Status)
	if err == sql.ErrNoRows {
		return seek, err, http.StatusNotFound
	} else if err != nil {
		return seek, err, http.StatusInternalServerError
	}

	return seek, nil, http.StatusOK
}

// Inserts the given seek (belonging to userId) into the database. Returns
// seek with its new KeyID added.
func CreateSeek(db *sql.DB, seek Seek, userId int) (Seek, error, int) {
	seek.UserID = userId

	// Insert seek
	stmt := psql.Insert("seeks").
		Columns("title", "description", "user_id", "saved_search_id",
			"notify_enabled", "status").
		Values(seek.Title, seek.Description, userId, seek.SavedSearchID,
			seek.NotifyEnabled, seek.Status).
		Suffix("RETURNING key_id, creation_date")

	// Query db for seek
	rows, err := stmt.RunWith(db).Query()
	if err != nil {
		return seek, err, http.StatusInternalServerError
	}
	defer rows.Close()

	// Populate seek struct
	rows.Next()
	err = rows.Scan(&seek.KeyID, &seek.CreationDate)
	if err != nil {
		return seek, err, http.StatusInternalServerError
	}

	return seek, nil, http.StatusCreated
}

// Overwrites the seek in the database with the given id with the given seek
// (belonging to userId). Returns the updated seek.
func UpdateSeek(db *sql.DB, id string, seek Seek, userId int) (error, int) {
	seek.UserID = userId

	// Update seek
	stmt := psql.Update("seeks").
		SetMap(map[string]interface{}{
			"title":           seek.Title,
			"description":     seek.Description,
			"saved_search_id": seek.SavedSearchID,
			"notify_enabled":  seek.NotifyEnabled}).
		Where(sq.Eq{"seeks.key_id": id,
			"seeks.user_id": userId})

	// Query db for seek
	result, err := stmt.RunWith(db).Exec()
	return getUpdateResultCode(result, err)
}

// Deletes the seek in the database with the given id with the given seek
// (belonging to userId).
func DeleteSeek(db *sql.DB, id string, userId int) (error, int) {
	// Update seek
	stmt := psql.Delete("seeks").
		Where(sq.Eq{"seeks.key_id": id,
			"seeks.user_id": userId})

	// Query db for seek
	result, err := stmt.RunWith(db).Exec()
	return getUpdateResultCode(result, err)
}
<|endoftext|>


<commit_before>package notify

import (
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/godbus/dbus/v5"
)

const (
	dbusRemoveMatch            = "org.freedesktop.DBus.RemoveMatch"
	dbusAddMatch               = "org.freedesktop.DBus.AddMatch"
	dbusObjectPath             = "/org/freedesktop/Notifications" // the DBUS object path
	dbusNotificationsInterface = "org.freedesktop.Notifications"  // DBUS Interface
	signalNotificationClosed   = "org.freedesktop.Notifications.NotificationClosed"
	signalActionInvoked        = "org.freedesktop.Notifications.ActionInvoked"
	callGetCapabilities        = "org.freedesktop.Notifications.GetCapabilities"
	callCloseNotification      = "org.freedesktop.Notifications.CloseNotification"
	callNotify                 = "org.freedesktop.Notifications.Notify"
	callGetServerInformation   = "org.freedesktop.Notifications.GetServerInformation"

	channelBufferSize = 10
)

// Notification holds all information needed for creating a notification
type Notification struct {
	AppName string
	// Setting ReplacesID atomically replaces the notification with this ID.
	// Optional.
	ReplacesID uint32
	// See predefined icons here: http://standards.freedesktop.org/icon-naming-spec/icon-naming-spec-latest.html
	// Optional.
	AppIcon string
	Summary string
	Body    string
	// Actions are tuples of (action_key, label), e.g.: []string{"cancel", "Cancel", "open", "Open"}
	Actions []string
	Hints   map[string]dbus.Variant
	// ExpireTimeout: milliseconds to show notification
	ExpireTimeout int32
}

// SendNotification is provided for convenience.
// Use if you only want to deliver a notification and dont care about events.
func SendNotification(conn *dbus.Conn, note Notification) (uint32, error) {
	actions := len(note.Actions)
	if (actions % 2) != 0 {
		return 0, errors.New("actions must be pairs of (key, label)")
	}

	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callNotify, 0,
		note.AppName,
		note.ReplacesID,
		note.AppIcon,
		note.Summary,
		note.Body,
		note.Actions,
		note.Hints,
		note.ExpireTimeout)
	if call.Err != nil {
		return 0, fmt.Errorf("error sending notification: %w", call.Err)
	}
	var ret uint32
	err := call.Store(&ret)
	if err != nil {
		return ret, fmt.Errorf("error getting uint32 ret value: %w", err)
	}
	return ret, nil
}

// ServerInformation is a holder for information returned by
// GetServerInformation call.
type ServerInformation struct {
	Name        string
	Vendor      string
	Version     string
	SpecVersion string
}

// GetServerInformation returns the information on the server.
//
// org.freedesktop.Notifications.GetServerInformation
//
//  GetServerInformation Return Values
//
//		Name		 Type	  Description
//		name		 STRING	  The product name of the server.
//		vendor		 STRING	  The vendor name. For example, "KDE," "GNOME," "freedesktop.org," or "Microsoft."
//		version		 STRING	  The server's version number.
//		spec_version STRING	  The specification version the server is compliant with.
//
func GetServerInformation(conn *dbus.Conn) (ServerInformation, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	if obj == nil {
		return ServerInformation{}, errors.New("error creating dbus call object")
	}
	call := obj.Call(callGetServerInformation, 0)
	if call.Err != nil {
		return ServerInformation{}, fmt.Errorf("error calling %v: %v", callGetServerInformation, call.Err)
	}

	ret := ServerInformation{}
	err := call.Store(&ret.Name, &ret.Vendor, &ret.Version, &ret.SpecVersion)
	if err != nil {
		return ret, fmt.Errorf("error reading %v return values: %v", callGetServerInformation, err)
	}
	return ret, nil
}

// GetCapabilities gets the capabilities of the notification server.
// This call takes no parameters.
// It returns an array of strings. Each string describes an optional capability implemented by the server.
//
// See also: https://developer.gnome.org/notification-spec/
// GetCapabilities provide an exported method for this operation
func GetCapabilities(conn *dbus.Conn) ([]string, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callGetCapabilities, 0)
	if call.Err != nil {
		log.Printf("error calling GetCapabilities: %v", call.Err)
		return []string{}, call.Err
	}
	var ret []string
	err := call.Store(&ret)
	if err != nil {
		log.Printf("error getting capabilities ret value: %v", err)
		return ret, err
	}
	return ret, nil
}

// Notifier is an interface for implementing the operations supported by the
// freedesktop DBus Notifications object.
//
// New() sets up a Notifier that listens on dbus' signals regarding
// Notifications: NotificationClosed and ActionInvoked.
//
// Note this also means the caller MUST consume output from these channels,
// given in methods NotificationClosed() and ActionInvoked().
// Users that only want to send a simple notification, but don't care about
// interactions, see exported method: SendNotification(conn, Notification)
//
// Caller is also responsible to call Close() before exiting,
// to shut down event loop and cleanup.
type Notifier interface {
	SendNotification(n Notification) (uint32, error)
	GetCapabilities() ([]string, error)
	GetServerInformation() (ServerInformation, error)
	CloseNotification(id int) (bool, error)
	// Deprecated: Use NotificationClosedHandler
	//NotificationClosed() <-chan *NotificationClosedSignal
	//ActionInvoked() <-chan *ActionInvokedSignal
	Close() error
}

// NotificationClosedHandler is called when we receive a NotificationClosed signal
type NotificationClosedHandler func(*NotificationClosedSignal)

// ActionInvokedHandler is called when we receive a signal that one of the actions passed was invoked.
//
// Note that invoking an action often also produces a NotificationClosedSignal,
// so you might receive both a Closed signal and a ActionInvoked signal.
// I suspect this detail is implementation specific for the UI interaction,
// and does at least happen on XFCE4.
type ActionInvokedHandler func(*ActionInvokedSignal)

// ActionInvokedSignal holds data from any signal received regarding Actions invoked
type ActionInvokedSignal struct {
	// ID of the Notification the action was invoked for
	ID uint32
	// Key from the tuple (action_key, label)
	ActionKey string
}

// notifier implements Notifier interface
type notifier struct {
	conn     *dbus.Conn
	signal   chan *dbus.Signal
	done     chan bool
	onClosed NotificationClosedHandler
	onAction ActionInvokedHandler
	wg       *sync.WaitGroup
	log      logger
}

type logger interface {
	Printf(format string, v ...interface{})
}

// Option overrides certain parts of a Notifier
type Option func(*notifier)

// WithLogger sets a new logger func
func WithLogger(logz logger) Option {
	return func(n *notifier) {
		n.log = logz
	}
}

// WithOnAction sets ActionInvokedHandler handler
func WithOnAction(h ActionInvokedHandler) Option {
	return func(n *notifier) {
		n.onAction = h
	}
}

// WithOnClosed sets NotificationClosed handler
func WithOnClosed(h NotificationClosedHandler) Option {
	return func(n *notifier) {
		n.onClosed = h
	}
}

// New creates a new Notifier using conn.
// See also: Notifier
func New(conn *dbus.Conn, opts ...Option) (Notifier, error) {
	n := &notifier{
		conn:     conn,
		signal:   make(chan *dbus.Signal, channelBufferSize),
		done:     make(chan bool),
		wg:       &sync.WaitGroup{},
		onClosed: func(s *NotificationClosedSignal) {},
		onAction: func(s *ActionInvokedSignal) {},
		log:      &loggerWrapper{"notify: "},
	}

	for _, val := range opts {
		val(n)
	}

	// add a listener in dbus for signals to Notification interface.
	err := n.conn.AddMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)
	if err != nil {
		return nil, err
	}

	// start eventloop
	go n.eventLoop()

	// register in dbus for signal delivery
	n.conn.Signal(n.signal)

	return n, nil
}

func (n notifier) eventLoop() {
	n.wg.Add(1)
	defer n.wg.Done()
	for {
		select {
		case signal := <-n.signal:
			n.handleSignal(signal)
		case <-n.done:
			n.log.Printf("Got Close() signal, shutting down...")
			return
		}
	}
}

// signal handler that translates and sends notifications to channels
func (n notifier) handleSignal(signal *dbus.Signal) {
	switch signal.Name {
	case signalNotificationClosed:
		nc := &NotificationClosedSignal{
			ID:     signal.Body[0].(uint32),
			Reason: Reason(signal.Body[1].(uint32)),
		}
		n.onClosed(nc)
	case signalActionInvoked:
		is := &ActionInvokedSignal{
			ID:        signal.Body[0].(uint32),
			ActionKey: signal.Body[1].(string),
		}
		n.onAction(is)
	default:
		n.log.Printf("Received unknown signal: %+v", signal)
	}
}

func (n *notifier) GetCapabilities() ([]string, error) {
	return GetCapabilities(n.conn)
}
func (n *notifier) GetServerInformation() (ServerInformation, error) {
	return GetServerInformation(n.conn)
}

// SendNotification sends a notification to the notification server and returns the ID or an error.
//
// Implements dbus call:
//
//     UINT32 org.freedesktop.Notifications.Notify (
//	       STRING app_name,
//	       UINT32 replaces_id,
//	       STRING app_icon,
//	       STRING summary,
//	       STRING body,
//	       ARRAY  actions,
//	       DICT   hints,
//	       INT32  expire_timeout
//     );
//
//		Name	    	Type	Description
//		app_name		STRING	The optional name of the application sending the notification. Can be blank.
//		replaces_id	    UINT32	The optional notification ID that this notification replaces. The server must atomically (ie with no flicker or other visual cues) replace the given notification with this one. This allows clients to effectively modify the notification while it's active. A value of value of 0 means that this notification won't replace any existing notifications.
//		app_icon		STRING	The optional program icon of the calling application. Can be an empty string, indicating no icon.
//		summary		    STRING	The summary text briefly describing the notification.
//		body			STRING	The optional detailed body text. Can be empty.
//		actions		    ARRAY	Actions are sent over as a list of pairs. Each even element in the list (starting at index 0) represents the identifier for the action. Each odd element in the list is the localized string that will be displayed to the user.
//		hints	        DICT	Optional hints that can be passed to the server from the client program. Although clients and servers should never assume each other supports any specific hints, they can be used to pass along information, such as the process PID or window ID, that the server may be able to make use of. See Hints. Can be empty.
//      expire_timeout  INT32   The timeout time in milliseconds since the display of the notification at which the notification should automatically close.
//								If -1, the notification's expiration time is dependent on the notification server's settings, and may vary for the type of notification. If 0, never expire.
//
// If replaces_id is 0, the return value is a UINT32 that represent the notification.
// It is unique, and will not be reused unless a MAXINT number of notifications have been generated.
//  An acceptable implementation may just use an incrementing counter for the ID.
// The returned ID is always greater than zero. Servers must make sure not to return zero as an ID.
//
// If replaces_id is not 0, the returned value is the same value as replaces_id.
func (n *notifier) SendNotification(note Notification) (uint32, error) {
	return SendNotification(n.conn, note)
}

// CloseNotification causes a notification to be forcefully closed and removed from the user's view.
// It can be used, for example, in the event that what the notification pertains to is no longer relevant,
// or to cancel a notification with no expiration time.
//
// The NotificationClosed (dbus) signal is emitted by this method.
// If the notification no longer exists, an empty D-BUS Error message is sent back.
func (n *notifier) CloseNotification(id int) (bool, error) {
	obj := n.conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callCloseNotification, 0, uint32(id))
	if call.Err != nil {
		return false, call.Err
	}
	return true, nil
}

// NotificationClosedSignal holds data for *Closed callbacks from Notifications Interface.
type NotificationClosedSignal struct {
	// ID of the Notification the signal was invoked for
	ID uint32
	// A reason given if known
	Reason Reason
}

// Reason for the closed notification
type Reason uint32

const (
	// ReasonExpired when a notification expired
	ReasonExpired Reason = 1

	// ReasonDismissedByUser when a notification has been dismissed by a user
	ReasonDismissedByUser Reason = 2

	// ReasonClosedByCall when a notification has been closed by a call to CloseNotification
	ReasonClosedByCall Reason = 3

	// ReasonUnknown when as notification has been closed for an unknown reason
	ReasonUnknown Reason = 4
)

func (r Reason) String() string {
	switch r {
	case ReasonExpired:
		return "Expired"
	case ReasonDismissedByUser:
		return "DismissedByUser"
	case ReasonClosedByCall:
		return "ClosedByCall"
	case ReasonUnknown:
		return "Unknown"
	default:
		return "Other"
	}
}

// Close cleans up and shuts down signal delivery loop
func (n *notifier) Close() error {
	n.done <- true

	// remove signal reception
	n.conn.RemoveSignal(n.signal)

	// unregister in dbus:
	errRemoveMatch := n.conn.RemoveMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)

	close(n.done)

	// wait for eventloop to shut down...
	n.wg.Wait()

	return errRemoveMatch
}

type loggerWrapper struct {
	prefix string
}

func (l *loggerWrapper) Printf(format string, v ...interface{}) {
	log.Printf(l.prefix+format, v...)
}

// NotificationClosed returns a receive only channel that sends
// NotificationClosedSignal for signals.
//
// The chan must be drained or event delivery will stall.
// Deprecated: Use NotificationClosedHandler
//func (n *notifier) NotificationClosed() <-chan *NotificationClosedSignal {
//	return n.closer
//}
<commit_msg>Update notification.go<commit_after>package notify

import (
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/godbus/dbus/v5"
)

const (
	dbusRemoveMatch            = "org.freedesktop.DBus.RemoveMatch"
	dbusAddMatch               = "org.freedesktop.DBus.AddMatch"
	dbusObjectPath             = "/org/freedesktop/Notifications" // the DBUS object path
	dbusNotificationsInterface = "org.freedesktop.Notifications"  // DBUS Interface
	signalNotificationClosed   = "org.freedesktop.Notifications.NotificationClosed"
	signalActionInvoked        = "org.freedesktop.Notifications.ActionInvoked"
	callGetCapabilities        = "org.freedesktop.Notifications.GetCapabilities"
	callCloseNotification      = "org.freedesktop.Notifications.CloseNotification"
	callNotify                 = "org.freedesktop.Notifications.Notify"
	callGetServerInformation   = "org.freedesktop.Notifications.GetServerInformation"

	channelBufferSize = 10
)

// Notification holds all information needed for creating a notification
type Notification struct {
	AppName string
	// Setting ReplacesID atomically replaces the notification with this ID.
	// Optional.
	ReplacesID uint32
	// See predefined icons here: http://standards.freedesktop.org/icon-naming-spec/icon-naming-spec-latest.html
	// Optional.
	AppIcon string
	Summary string
	Body    string
	// Actions are tuples of (action_key, label), e.g.: []string{"cancel", "Cancel", "open", "Open"}
	Actions []string
	Hints   map[string]dbus.Variant
	// ExpireTimeout: milliseconds to show notification
	ExpireTimeout int32
}

// SendNotification is provided for convenience.
// Use if you only want to deliver a notification and dont care about events.
func SendNotification(conn *dbus.Conn, note Notification) (uint32, error) {
	actions := len(note.Actions)
	if (actions % 2) != 0 {
		return 0, errors.New("actions must be pairs of (key, label)")
	}

	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callNotify, 0,
		note.AppName,
		note.ReplacesID,
		note.AppIcon,
		note.Summary,
		note.Body,
		note.Actions,
		note.Hints,
		note.ExpireTimeout)
	if call.Err != nil {
		return 0, fmt.Errorf("error sending notification: %w", call.Err)
	}
	var ret uint32
	err := call.Store(&ret)
	if err != nil {
		return ret, fmt.Errorf("error getting uint32 ret value: %w", err)
	}
	return ret, nil
}

// ServerInformation is a holder for information returned by
// GetServerInformation call.
type ServerInformation struct {
	Name        string
	Vendor      string
	Version     string
	SpecVersion string
}

// GetServerInformation returns the information on the server.
//
// org.freedesktop.Notifications.GetServerInformation
//
//  GetServerInformation Return Values
//
//		Name		 Type	  Description
//		name		 STRING	  The product name of the server.
//		vendor		 STRING	  The vendor name. For example, "KDE," "GNOME," "freedesktop.org," or "Microsoft."
//		version		 STRING	  The server's version number.
//		spec_version STRING	  The specification version the server is compliant with.
//
func GetServerInformation(conn *dbus.Conn) (ServerInformation, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	if obj == nil {
		return ServerInformation{}, errors.New("error creating dbus call object")
	}
	call := obj.Call(callGetServerInformation, 0)
	if call.Err != nil {
		return ServerInformation{}, fmt.Errorf("error calling %v: %v", callGetServerInformation, call.Err)
	}

	ret := ServerInformation{}
	err := call.Store(&ret.Name, &ret.Vendor, &ret.Version, &ret.SpecVersion)
	if err != nil {
		return ret, fmt.Errorf("error reading %v return values: %v", callGetServerInformation, err)
	}
	return ret, nil
}

// GetCapabilities gets the capabilities of the notification server.
// This call takes no parameters.
// It returns an array of strings. Each string describes an optional capability implemented by the server.
//
// See also: https://developer.gnome.org/notification-spec/
// GetCapabilities provide an exported method for this operation
func GetCapabilities(conn *dbus.Conn) ([]string, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callGetCapabilities, 0)
	if call.Err != nil {
		log.Printf("error calling GetCapabilities: %v", call.Err)
		return []string{}, call.Err
	}
	var ret []string
	err := call.Store(&ret)
	if err != nil {
		log.Printf("error getting capabilities ret value: %v", err)
		return ret, err
	}
	return ret, nil
}

// Notifier is an interface for implementing the operations supported by the
// freedesktop DBus Notifications object.
//
// New() sets up a Notifier that listens on dbus' signals regarding
// Notifications: NotificationClosed and ActionInvoked.
//
// Note this also means the caller MUST consume output from these channels,
// given in methods NotificationClosed() and ActionInvoked().
// Users that only want to send a simple notification, but don't care about
// interactions, see exported method: SendNotification(conn, Notification)
//
// Caller is also responsible to call Close() before exiting,
// to shut down event loop and cleanup.
type Notifier interface {
	SendNotification(n Notification) (uint32, error)
	GetCapabilities() ([]string, error)
	GetServerInformation() (ServerInformation, error)
	CloseNotification(id int) (bool, error)
	// Deprecated: Use NotificationClosedHandler
	//NotificationClosed() <-chan *NotificationClosedSignal
	//ActionInvoked() <-chan *ActionInvokedSignal
	Close() error
}

// NotificationClosedHandler is called when we receive a NotificationClosed signal
type NotificationClosedHandler func(*NotificationClosedSignal)

// ActionInvokedHandler is called when we receive a signal that one of the actions passed was invoked.
//
// Note that invoking an action often also produces a NotificationClosedSignal,
// so you might receive both a Closed signal and a ActionInvoked signal.
// I suspect this detail is implementation specific for the UI interaction,
// and does at least happen on XFCE4.
type ActionInvokedHandler func(*ActionInvokedSignal)

// ActionInvokedSignal holds data from any signal received regarding Actions invoked
type ActionInvokedSignal struct {
	// ID of the Notification the action was invoked for
	ID uint32
	// Key from the tuple (action_key, label)
	ActionKey string
}

// notifier implements Notifier interface
type notifier struct {
	conn     *dbus.Conn
	signal   chan *dbus.Signal
	done     chan bool
	onClosed NotificationClosedHandler
	onAction ActionInvokedHandler
	wg       *sync.WaitGroup
	log      logger
}

type logger interface {
	Printf(format string, v ...interface{})
}

// Option overrides certain parts of a Notifier
type Option func(*notifier)

// WithLogger sets a new logger func
func WithLogger(logz logger) Option {
	return func(n *notifier) {
		n.log = logz
	}
}

// WithOnAction sets ActionInvokedHandler handler
func WithOnAction(h ActionInvokedHandler) Option {
	return func(n *notifier) {
		n.onAction = h
	}
}

// WithOnClosed sets NotificationClosed handler
func WithOnClosed(h NotificationClosedHandler) Option {
	return func(n *notifier) {
		n.onClosed = h
	}
}

// New creates a new Notifier using conn.
// See also: Notifier
func New(conn *dbus.Conn, opts ...Option) (Notifier, error) {
	n := &notifier{
		conn:     conn,
		signal:   make(chan *dbus.Signal, channelBufferSize),
		done:     make(chan bool),
		wg:       &sync.WaitGroup{},
		onClosed: func(s *NotificationClosedSignal) {},
		onAction: func(s *ActionInvokedSignal) {},
		log:      &loggerWrapper{"notify: "},
	}

	for _, val := range opts {
		val(n)
	}

	// add a listener in dbus for signals to Notification interface.
	err := n.conn.AddMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)
	if err != nil {
		return nil, err
	}

	// start eventloop
	go n.eventLoop()

	// register in dbus for signal delivery
	n.conn.Signal(n.signal)

	return n, nil
}

func (n notifier) eventLoop() {
	n.wg.Add(1)
	defer n.wg.Done()
	for {
		select {
		case signal := <-n.signal:
			n.handleSignal(signal)
		case <-n.done:
			n.log.Printf("Got Close() signal, shutting down...")
			return
		}
	}
}

// signal handler that translates and sends notifications to channels
func (n notifier) handleSignal(signal *dbus.Signal) {
	switch signal.Name {
	case signalNotificationClosed:
		nc := &NotificationClosedSignal{
			ID:     signal.Body[0].(uint32),
			Reason: Reason(signal.Body[1].(uint32)),
		}
		n.onClosed(nc)
	case signalActionInvoked:
		is := &ActionInvokedSignal{
			ID:        signal.Body[0].(uint32),
			ActionKey: signal.Body[1].(string),
		}
		n.onAction(is)
	default:
		n.log.Printf("Received unknown signal: %+v", signal)
	}
}

func (n *notifier) GetCapabilities() ([]string, error) {
	return GetCapabilities(n.conn)
}
func (n *notifier) GetServerInformation() (ServerInformation, error) {
	return GetServerInformation(n.conn)
}

// SendNotification sends a notification to the notification server and returns the ID or an error.
//
// Implements dbus call:
//
//     UINT32 org.freedesktop.Notifications.Notify (
//	       STRING app_name,
//	       UINT32 replaces_id,
//	       STRING app_icon,
//	       STRING summary,
//	       STRING body,
//	       ARRAY  actions,
//	       DICT   hints,
//	       INT32  expire_timeout
//     );
//
//		Name	    	Type	Description
//		app_name		STRING	The optional name of the application sending the notification. Can be blank.
//		replaces_id	    UINT32	The optional notification ID that this notification replaces. The server must atomically (ie with no flicker or other visual cues) replace the given notification with this one. This allows clients to effectively modify the notification while it's active. A value of value of 0 means that this notification won't replace any existing notifications.
//		app_icon		STRING	The optional program icon of the calling application. Can be an empty string, indicating no icon.
//		summary		    STRING	The summary text briefly describing the notification.
//		body			STRING	The optional detailed body text. Can be empty.
//		actions		    ARRAY	Actions are sent over as a list of pairs. Each even element in the list (starting at index 0) represents the identifier for the action. Each odd element in the list is the localized string that will be displayed to the user.
//		hints	        DICT	Optional hints that can be passed to the server from the client program. Although clients and servers should never assume each other supports any specific hints, they can be used to pass along information, such as the process PID or window ID, that the server may be able to make use of. See Hints. Can be empty.
//      expire_timeout  INT32   The timeout time in milliseconds since the display of the notification at which the notification should automatically close.
//								If -1, the notification's expiration time is dependent on the notification server's settings, and may vary for the type of notification. If 0, never expire.
//
// If replaces_id is 0, the return value is a UINT32 that represent the notification.
// It is unique, and will not be reused unless a MAXINT number of notifications have been generated.
// An acceptable implementation may just use an incrementing counter for the ID.
// The returned ID is always greater than zero. Servers must make sure not to return zero as an ID.
//
// If replaces_id is not 0, the returned value is the same value as replaces_id.
func (n *notifier) SendNotification(note Notification) (uint32, error) {
	return SendNotification(n.conn, note)
}

// CloseNotification causes a notification to be forcefully closed and removed from the user's view.
// It can be used, for example, in the event that what the notification pertains to is no longer relevant,
// or to cancel a notification with no expiration time.
//
// The NotificationClosed (dbus) signal is emitted by this method.
// If the notification no longer exists, an empty D-BUS Error message is sent back.
func (n *notifier) CloseNotification(id int) (bool, error) {
	obj := n.conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callCloseNotification, 0, uint32(id))
	if call.Err != nil {
		return false, call.Err
	}
	return true, nil
}

// NotificationClosedSignal holds data for *Closed callbacks from Notifications Interface.
type NotificationClosedSignal struct {
	// ID of the Notification the signal was invoked for
	ID uint32
	// A reason given if known
	Reason Reason
}

// Reason for the closed notification
type Reason uint32

const (
	// ReasonExpired when a notification expired
	ReasonExpired Reason = 1

	// ReasonDismissedByUser when a notification has been dismissed by a user
	ReasonDismissedByUser Reason = 2

	// ReasonClosedByCall when a notification has been closed by a call to CloseNotification
	ReasonClosedByCall Reason = 3

	// ReasonUnknown when as notification has been closed for an unknown reason
	ReasonUnknown Reason = 4
)

func (r Reason) String() string {
	switch r {
	case ReasonExpired:
		return "Expired"
	case ReasonDismissedByUser:
		return "DismissedByUser"
	case ReasonClosedByCall:
		return "ClosedByCall"
	case ReasonUnknown:
		return "Unknown"
	default:
		return "Other"
	}
}

// Close cleans up and shuts down signal delivery loop
func (n *notifier) Close() error {
	n.done <- true

	// remove signal reception
	n.conn.RemoveSignal(n.signal)

	// unregister in dbus:
	errRemoveMatch := n.conn.RemoveMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)

	close(n.done)

	// wait for eventloop to shut down...
	n.wg.Wait()

	return errRemoveMatch
}

type loggerWrapper struct {
	prefix string
}

func (l *loggerWrapper) Printf(format string, v ...interface{}) {
	log.Printf(l.prefix+format, v...)
}

// NotificationClosed returns a receive only channel that sends
// NotificationClosedSignal for signals.
//
// The chan must be drained or event delivery will stall.
// Deprecated: Use NotificationClosedHandler
//func (n *notifier) NotificationClosed() <-chan *NotificationClosedSignal {
//	return n.closer
//}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	errs "errors"

	"go-micro.dev/v4/codec"
	raw "go-micro.dev/v4/codec/bytes"
	"go-micro.dev/v4/codec/grpc"
	"go-micro.dev/v4/codec/json"
	"go-micro.dev/v4/codec/jsonrpc"
	"go-micro.dev/v4/codec/proto"
	"go-micro.dev/v4/codec/protorpc"
	"go-micro.dev/v4/errors"
	"go-micro.dev/v4/registry"
	"go-micro.dev/v4/transport"
)

const (
	lastStreamResponseError = "EOS"
)

// serverError represents an error that has been returned from
// the remote side of the RPC connection.
type serverError string

func (e serverError) Error() string {
	return string(e)
}

// errShutdown holds the specific error for closing/closed connections
var (
	errShutdown = errs.New("connection is shut down")
)

type rpcCodec struct {
	client transport.Client
	codec  codec.Codec

	req *transport.Message
	buf *readWriteCloser

	// signify if its a stream
	stream string
}

type readWriteCloser struct {
	wbuf *bytes.Buffer
	rbuf *bytes.Buffer
}

var (
	DefaultContentType = "application/protobuf"

	DefaultCodecs = map[string]codec.NewCodec{
		"application/grpc":         grpc.NewCodec,
		"application/grpc+json":    grpc.NewCodec,
		"application/grpc+proto":   grpc.NewCodec,
		"application/protobuf":     proto.NewCodec,
		"application/json":         json.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": raw.NewCodec,
	}

	// TODO: remove legacy codec list
	defaultCodecs = map[string]codec.NewCodec{
		"application/json":         jsonrpc.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/protobuf":     protorpc.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": protorpc.NewCodec,
	}
)

func (rwc *readWriteCloser) Read(p []byte) (n int, err error) {
	return rwc.rbuf.Read(p)
}

func (rwc *readWriteCloser) Write(p []byte) (n int, err error) {
	return rwc.wbuf.Write(p)
}

func (rwc *readWriteCloser) Close() error {
	rwc.rbuf.Reset()
	rwc.wbuf.Reset()
	return nil
}

func getHeaders(m *codec.Message) {
	set := func(v, hdr string) string {
		if len(v) > 0 {
			return v
		}
		return m.Header[hdr]
	}

	// check error in header
	m.Error = set(m.Error, "Micro-Error")

	// check endpoint in header
	m.Endpoint = set(m.Endpoint, "Micro-Endpoint")

	// check method in header
	m.Method = set(m.Method, "Micro-Method")

	// set the request id
	m.Id = set(m.Id, "Micro-Id")
}

func setHeaders(m *codec.Message, stream string) {
	set := func(hdr, v string) {
		if len(v) == 0 {
			return
		}
		m.Header[hdr] = v
	}

	set("Micro-Id", m.Id)
	set("Micro-Service", m.Target)
	set("Micro-Method", m.Method)
	set("Micro-Endpoint", m.Endpoint)
	set("Micro-Error", m.Error)

	if len(stream) > 0 {
		set("Micro-Stream", stream)
	}
}

// setupProtocol sets up the old protocol
func setupProtocol(msg *transport.Message, node *registry.Node) codec.NewCodec {
	protocol := node.Metadata["protocol"]

	// got protocol
	if len(protocol) > 0 {
		return nil
	}

	// processing topic publishing
	if len(msg.Header["Micro-Topic"]) > 0 {
		return nil
	}

	// no protocol use old codecs
	switch msg.Header["Content-Type"] {
	case "application/json":
		msg.Header["Content-Type"] = "application/json-rpc"
	case "application/protobuf":
		msg.Header["Content-Type"] = "application/proto-rpc"
	}

	// now return codec
	return defaultCodecs[msg.Header["Content-Type"]]
}

func newRpcCodec(req *transport.Message, client transport.Client, c codec.NewCodec, stream string) codec.Codec {
	rwc := &readWriteCloser{
		wbuf: bytes.NewBuffer(nil),
		rbuf: bytes.NewBuffer(nil),
	}
	r := &rpcCodec{
		buf:    rwc,
		client: client,
		codec:  c(rwc),
		req:    req,
		stream: stream,
	}
	return r
}

func (c *rpcCodec) Write(m *codec.Message, body interface{}) error {
	c.buf.wbuf.Reset()

	// create header
	if m.Header == nil {
		m.Header = map[string]string{}
	}

	// copy original header
	for k, v := range c.req.Header {
		m.Header[k] = v
	}

	// set the mucp headers
	setHeaders(m, c.stream)

	// if body is bytes Frame don't encode
	if body != nil {
		if b, ok := body.(*raw.Frame); ok {
			// set body
			m.Body = b.Data
		} else {
			// write to codec
			if err := c.codec.Write(m, body); err != nil {
				return errors.InternalServerError("go.micro.client.codec", err.Error())
			}
			// set body
			m.Body = c.buf.wbuf.Bytes()
		}
	}

	// create new transport message
	msg := transport.Message{
		Header: m.Header,
		Body:   m.Body,
	}

	// send the request
	if err := c.client.Send(&msg); err != nil {
		return errors.InternalServerError("go.micro.client.transport", err.Error())
	}

	return nil
}

func (c *rpcCodec) ReadHeader(m *codec.Message, r codec.MessageType) error {
	var tm transport.Message

	// read message from transport
	if err := c.client.Recv(&tm); err != nil {
		return errors.InternalServerError("go.micro.client.transport", err.Error())
	}

	c.buf.rbuf.Reset()
	c.buf.rbuf.Write(tm.Body)

	// set headers from transport
	m.Header = tm.Header

	// read header
	err := c.codec.ReadHeader(m, r)

	// get headers
	getHeaders(m)

	// return header error
	if err != nil {
		return errors.InternalServerError("go.micro.client.codec", err.Error())
	}

	return nil
}

func (c *rpcCodec) ReadBody(b interface{}) error {
	// read body
	// read raw data
	if v, ok := b.(*raw.Frame); ok {
		v.Data = c.buf.rbuf.Bytes()
		return nil
	}

	if err := c.codec.ReadBody(b); err != nil {
		return errors.InternalServerError("go.micro.client.codec", err.Error())
	}
	return nil
}

func (c *rpcCodec) Close() error {
	c.buf.Close()
	c.codec.Close()
	if err := c.client.Close(); err != nil {
		return errors.InternalServerError("go.micro.client.transport", err.Error())
	}
	return nil
}

func (c *rpcCodec) String() string {
	return "rpc"
}
<commit_msg>default the content type to json (#2412)<commit_after>package client

import (
	"bytes"
	errs "errors"

	"go-micro.dev/v4/codec"
	raw "go-micro.dev/v4/codec/bytes"
	"go-micro.dev/v4/codec/grpc"
	"go-micro.dev/v4/codec/json"
	"go-micro.dev/v4/codec/jsonrpc"
	"go-micro.dev/v4/codec/proto"
	"go-micro.dev/v4/codec/protorpc"
	"go-micro.dev/v4/errors"
	"go-micro.dev/v4/registry"
	"go-micro.dev/v4/transport"
)

const (
	lastStreamResponseError = "EOS"
)

// serverError represents an error that has been returned from
// the remote side of the RPC connection.
type serverError string

func (e serverError) Error() string {
	return string(e)
}

// errShutdown holds the specific error for closing/closed connections
var (
	errShutdown = errs.New("connection is shut down")
)

type rpcCodec struct {
	client transport.Client
	codec  codec.Codec

	req *transport.Message
	buf *readWriteCloser

	// signify if its a stream
	stream string
}

type readWriteCloser struct {
	wbuf *bytes.Buffer
	rbuf *bytes.Buffer
}

var (
	DefaultContentType = "application/json"

	DefaultCodecs = map[string]codec.NewCodec{
		"application/grpc":         grpc.NewCodec,
		"application/grpc+json":    grpc.NewCodec,
		"application/grpc+proto":   grpc.NewCodec,
		"application/protobuf":     proto.NewCodec,
		"application/json":         json.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": raw.NewCodec,
	}

	// TODO: remove legacy codec list
	defaultCodecs = map[string]codec.NewCodec{
		"application/json":         jsonrpc.NewCodec,
		"application/json-rpc":     jsonrpc.NewCodec,
		"application/protobuf":     protorpc.NewCodec,
		"application/proto-rpc":    protorpc.NewCodec,
		"application/octet-stream": protorpc.NewCodec,
	}
)

func (rwc *readWriteCloser) Read(p []byte) (n int, err error) {
	return rwc.rbuf.Read(p)
}

func (rwc *readWriteCloser) Write(p []byte) (n int, err error) {
	return rwc.wbuf.Write(p)
}

func (rwc *readWriteCloser) Close() error {
	rwc.rbuf.Reset()
	rwc.wbuf.Reset()
	return nil
}

func getHeaders(m *codec.Message) {
	set := func(v, hdr string) string {
		if len(v) > 0 {
			return v
		}
		return m.Header[hdr]
	}

	// check error in header
	m.Error = set(m.Error, "Micro-Error")

	// check endpoint in header
	m.Endpoint = set(m.Endpoint, "Micro-Endpoint")

	// check method in header
	m.Method = set(m.Method, "Micro-Method")

	// set the request id
	m.Id = set(m.Id, "Micro-Id")
}

func setHeaders(m *codec.Message, stream string) {
	set := func(hdr, v string) {
		if len(v) == 0 {
			return
		}
		m.Header[hdr] = v
	}

	set("Micro-Id", m.Id)
	set("Micro-Service", m.Target)
	set("Micro-Method", m.Method)
	set("Micro-Endpoint", m.Endpoint)
	set("Micro-Error", m.Error)

	if len(stream) > 0 {
		set("Micro-Stream", stream)
	}
}

// setupProtocol sets up the old protocol
func setupProtocol(msg *transport.Message, node *registry.Node) codec.NewCodec {
	protocol := node.Metadata["protocol"]

	// got protocol
	if len(protocol) > 0 {
		return nil
	}

	// processing topic publishing
	if len(msg.Header["Micro-Topic"]) > 0 {
		return nil
	}

	// no protocol use old codecs
	switch msg.Header["Content-Type"] {
	case "application/json":
		msg.Header["Content-Type"] = "application/json-rpc"
	case "application/protobuf":
		msg.Header["Content-Type"] = "application/proto-rpc"
	}

	// now return codec
	return defaultCodecs[msg.Header["Content-Type"]]
}

func newRpcCodec(req *transport.Message, client transport.Client, c codec.NewCodec, stream string) codec.Codec {
	rwc := &readWriteCloser{
		wbuf: bytes.NewBuffer(nil),
		rbuf: bytes.NewBuffer(nil),
	}
	r := &rpcCodec{
		buf:    rwc,
		client: client,
		codec:  c(rwc),
		req:    req,
		stream: stream,
	}
	return r
}

func (c *rpcCodec) Write(m *codec.Message, body interface{}) error {
	c.buf.wbuf.Reset()

	// create header
	if m.Header == nil {
		m.Header = map[string]string{}
	}

	// copy original header
	for k, v := range c.req.Header {
		m.Header[k] = v
	}

	// set the mucp headers
	setHeaders(m, c.stream)

	// if body is bytes Frame don't encode
	if body != nil {
		if b, ok := body.(*raw.Frame); ok {
			// set body
			m.Body = b.Data
		} else {
			// write to codec
			if err := c.codec.Write(m, body); err != nil {
				return errors.InternalServerError("go.micro.client.codec", err.Error())
			}
			// set body
			m.Body = c.buf.wbuf.Bytes()
		}
	}

	// create new transport message
	msg := transport.Message{
		Header: m.Header,
		Body:   m.Body,
	}

	// send the request
	if err := c.client.Send(&msg); err != nil {
		return errors.InternalServerError("go.micro.client.transport", err.Error())
	}

	return nil
}

func (c *rpcCodec) ReadHeader(m *codec.Message, r codec.MessageType) error {
	var tm transport.Message

	// read message from transport
	if err := c.client.Recv(&tm); err != nil {
		return errors.InternalServerError("go.micro.client.transport", err.Error())
	}

	c.buf.rbuf.Reset()
	c.buf.rbuf.Write(tm.Body)

	// set headers from transport
	m.Header = tm.Header

	// read header
	err := c.codec.ReadHeader(m, r)

	// get headers
	getHeaders(m)

	// return header error
	if err != nil {
		return errors.InternalServerError("go.micro.client.codec", err.Error())
	}

	return nil
}

func (c *rpcCodec) ReadBody(b interface{}) error {
	// read body
	// read raw data
	if v, ok := b.(*raw.Frame); ok {
		v.Data = c.buf.rbuf.Bytes()
		return nil
	}

	if err := c.codec.ReadBody(b); err != nil {
		return errors.InternalServerError("go.micro.client.codec", err.Error())
	}
	return nil
}

func (c *rpcCodec) Close() error {
	c.buf.Close()
	c.codec.Close()
	if err := c.client.Close(); err != nil {
		return errors.InternalServerError("go.micro.client.transport", err.Error())
	}
	return nil
}

func (c *rpcCodec) String() string {
	return "rpc"
}
<|endoftext|>