<commit_before>package frames

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"time"
)

type ChannelDialer interface {
	io.Closer
	Dial() (net.Conn, error)
}

type queueResult struct {
	conn net.Conn
	err  error
}

type frameClient struct {
	c         net.Conn
	channels  map[uint16]*clientChannel
	egress    chan *FramePacket
	connqueue chan chan queueResult
}

func (fc *frameClient) handleOpened(pkt *FramePacket) {
	var opening chan queueResult
	select {
	case opening = <-fc.connqueue:
	default:
		log.Panicf("Opening response, but nobody's opening")
		return
	}

	if pkt.Status != FrameSuccess {
		err := frameError(*pkt)
		opening <- queueResult{err: err}
		return
	}

	fc.channels[pkt.Channel] = &clientChannel{
		fc,
		pkt.Channel,
		make(chan []byte, 8192),
		nil,
	}
	opening <- queueResult{fc.channels[pkt.Channel], nil}
}

func (fc *frameClient) handleClosed(pkt *FramePacket) {
	log.Panicf("Closing channel %v (unhandled)", pkt.Channel)
}

func (fc *frameClient) handleData(pkt *FramePacket) {
	ch := fc.channels[pkt.Channel]
	if ch == nil {
		log.Panicf("Data on closed channel: %v", pkt)
	}
	ch.incoming <- pkt.Data
}

func (fc *frameClient) readResponses() {
	defer fc.Close()
	for {
		hdr := make([]byte, minPktLen)
		_, err := io.ReadFull(fc.c, hdr)
		if err != nil {
			log.Printf("Error reading pkt header: %v", err)
			return
		}
		pkt := PacketFromHeader(hdr)
		_, err = io.ReadFull(fc.c, pkt.Data)
		if err != nil {
			log.Printf("Error reading pkt body: %v", err)
			return
		}

		switch pkt.Cmd {
		case FrameOpen:
			fc.handleOpened(&pkt)
		case FrameClose:
			fc.handleClosed(&pkt)
		case FrameData:
			fc.handleData(&pkt)
		default:
			panic("unhandled msg")
		}
	}
}

func (fc *frameClient) writeRequests() {
	for {
		e := <-fc.egress
		_, err := fc.c.Write(e.Bytes())
		if err != nil {
			panic(err)
		}
	}
}

// Convert a socket into a channel dialer.
func NewClient(c net.Conn) ChannelDialer {
	fc := &frameClient{
		c,
		map[uint16]*clientChannel{},
		make(chan *FramePacket, 256),
		make(chan chan queueResult, 16),
	}

	go fc.readResponses()
	go fc.writeRequests()

	return fc
}

func (f *frameClient) Close() error {
	for _, fc := range f.channels {
		fc.terminate()
	}
	return f.c.Close()
}

func (f *frameClient) Dial() (net.Conn, error) {
	pkt := &FramePacket{Cmd: FrameOpen}

	ch := make(chan queueResult)

	f.connqueue <- ch
	f.egress <- pkt

	qr := <-ch
	return qr.conn, qr.err
}

type clientChannel struct {
	fc       *frameClient
	channel  uint16
	incoming chan []byte
	current  []byte
}

func (f *clientChannel) Read(b []byte) (n int, err error) {
	if f.incoming == nil {
		return 0, io.EOF
	}
	read := 0
	for len(b) > 0 {
		if f.current == nil || len(f.current) == 0 {
			if read == 0 {
				f.current = <-f.incoming
			} else {
				var ok bool
				select {
				case f.current, ok = <-f.incoming:
					if !ok {
						return read, io.EOF
					}
				default:
					return read, nil
				}
			}
		}
		copied := copy(b, f.current)
		read += copied
		f.current = f.current[copied:]
		b = b[copied:]
	}
	return read, nil
}

func (f *clientChannel) Write(b []byte) (n int, err error) {
	if f.incoming == nil {
		return 0, io.EOF
	}

	bc := make([]byte, len(b))
	copy(bc, b)
	pkt := &FramePacket{
		Cmd:     FrameData,
		Channel: f.channel,
		Data:    bc,
	}
	f.fc.egress <- pkt
	return len(b), nil
}

func (f *clientChannel) Close() error {
	f.fc.egress <- &FramePacket{
		Cmd:     FrameClose,
		Channel: f.channel,
	}
	f.terminate()
	return nil
}

func (f *clientChannel) terminate() {
	close(f.incoming)
	f.incoming = nil
}

type frameAddr struct {
	a  net.Addr
	ch uint16
}

func (f frameAddr) Network() string {
	return f.a.String()
}

func (f frameAddr) String() string {
	return fmt.Sprintf("%v#%v", f.Network(), f.ch)
}

func (f *clientChannel) LocalAddr() net.Addr {
	return frameAddr{f.fc.c.LocalAddr(), f.channel}
}

func (f *clientChannel) RemoteAddr() net.Addr {
	return frameAddr{f.fc.c.RemoteAddr(), f.channel}
}

func (f *clientChannel) SetDeadline(t time.Time) error {
	return errors.New("Not Implemented")
}

func (f *clientChannel) SetReadDeadline(t time.Time) error {
	return errors.New("Not Implemented")
}

func (f *clientChannel) SetWriteDeadline(t time.Time) error {
	return errors.New("Not Implemented")
}

func (f *clientChannel) String() string {
	return fmt.Sprintf("ClientChannel{%v -> %v #%v}",
		f.fc.c.LocalAddr(), f.fc.c.RemoteAddr(), f.channel)
}
<commit_msg>Close client connections on cleanup.<commit_after>package frames

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"time"
)

type ChannelDialer interface {
	io.Closer
	Dial() (net.Conn, error)
}

type queueResult struct {
	conn net.Conn
	err  error
}

type frameClient struct {
	c         net.Conn
	channels  map[uint16]*clientChannel
	egress    chan *FramePacket
	connqueue chan chan queueResult
}

func (fc *frameClient) handleOpened(pkt *FramePacket) {
	var opening chan queueResult
	select {
	case opening = <-fc.connqueue:
	default:
		log.Panicf("Opening response, but nobody's opening")
		return
	}

	if pkt.Status != FrameSuccess {
		err := frameError(*pkt)
		opening <- queueResult{err: err}
		return
	}

	fc.channels[pkt.Channel] = &clientChannel{
		fc,
		pkt.Channel,
		make(chan []byte, 8192),
		nil,
	}
	opening <- queueResult{fc.channels[pkt.Channel], nil}
}

func (fc *frameClient) handleClosed(pkt *FramePacket) {
	log.Panicf("Closing channel %v (unhandled)", pkt.Channel)
}

func (fc *frameClient) handleData(pkt *FramePacket) {
	ch := fc.channels[pkt.Channel]
	if ch == nil {
		log.Panicf("Data on closed channel: %v", pkt)
	}
	ch.incoming <- pkt.Data
}

func (fc *frameClient) readResponses() {
	defer fc.Close()
	for {
		hdr := make([]byte, minPktLen)
		_, err := io.ReadFull(fc.c, hdr)
		if err != nil {
			log.Printf("Error reading pkt header: %v", err)
			return
		}
		pkt := PacketFromHeader(hdr)
		_, err = io.ReadFull(fc.c, pkt.Data)
		if err != nil {
			log.Printf("Error reading pkt body: %v", err)
			return
		}

		switch pkt.Cmd {
		case FrameOpen:
			fc.handleOpened(&pkt)
		case FrameClose:
			fc.handleClosed(&pkt)
		case FrameData:
			fc.handleData(&pkt)
		default:
			panic("unhandled msg")
		}
	}
}

func (fc *frameClient) writeRequests() {
	for {
		e := <-fc.egress
		_, err := fc.c.Write(e.Bytes())
		// Clean up on close
		if e.Cmd == FrameClose {
			delete(fc.channels, e.Channel)
		}
		if err != nil {
			panic(err)
		}
	}
}

// Convert a socket into a channel dialer.
func NewClient(c net.Conn) ChannelDialer {
	fc := &frameClient{
		c,
		map[uint16]*clientChannel{},
		make(chan *FramePacket, 256),
		make(chan chan queueResult, 16),
	}

	go fc.readResponses()
	go fc.writeRequests()

	return fc
}

func (f *frameClient) Close() error {
	for _, fc := range f.channels {
		fc.terminate()
	}
	return f.c.Close()
}

func (f *frameClient) Dial() (net.Conn, error) {
	pkt := &FramePacket{Cmd: FrameOpen}

	ch := make(chan queueResult)

	f.connqueue <- ch
	f.egress <- pkt

	qr := <-ch
	return qr.conn, qr.err
}

type clientChannel struct {
	fc       *frameClient
	channel  uint16
	incoming chan []byte
	current  []byte
}

func (f *clientChannel) Read(b []byte) (n int, err error) {
	if f.incoming == nil {
		return 0, io.EOF
	}
	read := 0
	for len(b) > 0 {
		if f.current == nil || len(f.current) == 0 {
			if read == 0 {
				f.current = <-f.incoming
			} else {
				var ok bool
				select {
				case f.current, ok = <-f.incoming:
					if !ok {
						return read, io.EOF
					}
				default:
					return read, nil
				}
			}
		}
		copied := copy(b, f.current)
		read += copied
		f.current = f.current[copied:]
		b = b[copied:]
	}
	return read, nil
}

func (f *clientChannel) Write(b []byte) (n int, err error) {
	if f.incoming == nil {
		return 0, io.EOF
	}

	bc := make([]byte, len(b))
	copy(bc, b)
	pkt := &FramePacket{
		Cmd:     FrameData,
		Channel: f.channel,
		Data:    bc,
	}
	f.fc.egress <- pkt
	return len(b), nil
}

func (f *clientChannel) Close() error {
	f.fc.egress <- &FramePacket{
		Cmd:     FrameClose,
		Channel: f.channel,
	}
	f.terminate()
	return nil
}

func (f *clientChannel) terminate() {
	close(f.incoming)
	f.incoming = nil
}

type frameAddr struct {
	a  net.Addr
	ch uint16
}

func (f frameAddr) Network() string {
	return f.a.String()
}

func (f frameAddr) String() string {
	return fmt.Sprintf("%v#%v", f.Network(), f.ch)
}

func (f *clientChannel) LocalAddr() net.Addr {
	return frameAddr{f.fc.c.LocalAddr(), f.channel}
}

func (f *clientChannel) RemoteAddr() net.Addr {
	return frameAddr{f.fc.c.RemoteAddr(), f.channel}
}

func (f *clientChannel) SetDeadline(t time.Time) error {
	return errors.New("Not Implemented")
}

func (f *clientChannel) SetReadDeadline(t time.Time) error {
	return errors.New("Not Implemented")
}

func (f *clientChannel) SetWriteDeadline(t time.Time) error {
	return errors.New("Not Implemented")
}

func (f *clientChannel) String() string {
	return fmt.Sprintf("ClientChannel{%v -> %v #%v}",
		f.fc.c.LocalAddr(), f.fc.c.RemoteAddr(), f.channel)
}
<|endoftext|>


<commit_before>package hdfs

import (
	hdfs "github.com/colinmarc/hdfs/protocol/hadoop_hdfs"
	"github.com/colinmarc/hdfs/rpc"
	"github.com/golang/protobuf/proto"
	"io"
	"io/ioutil"
	"os"
	"os/user"
)

// A Client represents a connection to an HDFS cluster
type Client struct {
	namenode *rpc.NamenodeConnection
	defaults *hdfs.FsServerDefaultsProto
}

// New returns a connected Client, or an error if it can't connect. The user
// will be the user the code is running under.
func New(address string) (*Client, error) {
	currentUser, err := user.Current()
	if err != nil {
		return nil, err
	}

	return NewForUser(address, currentUser.Username)
}

// NewForUser returns a connected Client with the user specified, or an error if
// it can't connect.
func NewForUser(address string, user string) (*Client, error) {
	namenode, err := rpc.NewNamenodeConnection(address, user)
	if err != nil {
		return nil, err
	}

	return &Client{namenode: namenode}, nil
}

// ReadFile reads the file named by filename and returns the contents.
func (c *Client) ReadFile(filename string) ([]byte, error) {
	f, err := c.Open(filename)
	if err != nil {
		return nil, err
	}

	return ioutil.ReadAll(f)
}

// CopyToLocal copies the HDFS file specified by src to the local file at dst.
// If dst already exists, it will be overwritten.
func (c *Client) CopyToLocal(src string, dst string) error {
	remote, err := c.Open(src)
	if err != nil {
		return err
	}

	local, err := os.Create(dst)
	if err != nil {
		return err
	}

	_, err = io.Copy(local, remote)
	return err
}

// CreateEmptyFile creates a empty file named by filename, with the permissions
// 0644.
func (c *Client) CreateEmptyFile(filename string) error {
	_, err := c.getFileInfo(filename)
	if err == nil {
		return &os.PathError{"create", filename, os.ErrExist}
	} else if !os.IsNotExist(err) {
		return &os.PathError{"create", filename, err}
	}

	defaults, err := c.fetchDefaults()
	if err != nil {
		return err
	}

	createReq := &hdfs.CreateRequestProto{
		Src:          proto.String(filename),
		Masked:       &hdfs.FsPermissionProto{Perm: proto.Uint32(uint32(0644))},
		ClientName:   proto.String(rpc.ClientName),
		CreateFlag:   proto.Uint32(1),
		CreateParent: proto.Bool(false),
		Replication:  proto.Uint32(defaults.GetReplication()),
		BlockSize:    proto.Uint64(defaults.GetBlockSize()),
	}
	createResp := &hdfs.CreateResponseProto{}

	err = c.namenode.Execute("create", createReq, createResp)
	if err != nil {
		if nnErr, ok := err.(*rpc.NamenodeError); ok {
			err = interpretException(nnErr.Exception, err)
		}

		return &os.PathError{"create", filename, err}
	}

	completeReq := &hdfs.CompleteRequestProto{
		Src:        proto.String(filename),
		ClientName: proto.String(rpc.ClientName),
	}
	completeResp := &hdfs.CompleteResponseProto{}

	err = c.namenode.Execute("complete", completeReq, completeResp)
	if err != nil {
		return &os.PathError{"create", filename, err}
	}

	return nil
}

func (c *Client) fetchDefaults() (*hdfs.FsServerDefaultsProto, error) {
	if c.defaults != nil {
		return c.defaults, nil
	}

	req := &hdfs.GetServerDefaultsRequestProto{}
	resp := &hdfs.GetServerDefaultsResponseProto{}

	err := c.namenode.Execute("getServerDefaults", req, resp)
	if err != nil {
		return nil, err
	}

	c.defaults = resp.GetServerDefaults()
	return c.defaults, nil
}
<commit_msg>Hack to workaround Current.user on linux<commit_after>package hdfs

import (
	hdfs "github.com/colinmarc/hdfs/protocol/hadoop_hdfs"
	"github.com/colinmarc/hdfs/rpc"
	"github.com/golang/protobuf/proto"
	"io"
	"io/ioutil"
	"os"
	"os/user"
)

// A Client represents a connection to an HDFS cluster
type Client struct {
	namenode *rpc.NamenodeConnection
	defaults *hdfs.FsServerDefaultsProto
}

// New returns a connected Client, or an error if it can't connect. The user
// will be the user the code is running under.
func New(address string) (*Client, error) {
	currentUser, err := user.Current()
	if err != nil {
		return nil, err
	}

	//	return NewForUser(address, currentUser.Username)
	return NewForUser(address, "nobody")
}

// NewForUser returns a connected Client with the user specified, or an error if
// it can't connect.
func NewForUser(address string, user string) (*Client, error) {
	namenode, err := rpc.NewNamenodeConnection(address, user)
	if err != nil {
		return nil, err
	}

	return &Client{namenode: namenode}, nil
}

// ReadFile reads the file named by filename and returns the contents.
func (c *Client) ReadFile(filename string) ([]byte, error) {
	f, err := c.Open(filename)
	if err != nil {
		return nil, err
	}

	return ioutil.ReadAll(f)
}

// CopyToLocal copies the HDFS file specified by src to the local file at dst.
// If dst already exists, it will be overwritten.
func (c *Client) CopyToLocal(src string, dst string) error {
	remote, err := c.Open(src)
	if err != nil {
		return err
	}

	local, err := os.Create(dst)
	if err != nil {
		return err
	}

	_, err = io.Copy(local, remote)
	return err
}

// CreateEmptyFile creates a empty file named by filename, with the permissions
// 0644.
func (c *Client) CreateEmptyFile(filename string) error {
	_, err := c.getFileInfo(filename)
	if err == nil {
		return &os.PathError{"create", filename, os.ErrExist}
	} else if !os.IsNotExist(err) {
		return &os.PathError{"create", filename, err}
	}

	defaults, err := c.fetchDefaults()
	if err != nil {
		return err
	}

	createReq := &hdfs.CreateRequestProto{
		Src:          proto.String(filename),
		Masked:       &hdfs.FsPermissionProto{Perm: proto.Uint32(uint32(0644))},
		ClientName:   proto.String(rpc.ClientName),
		CreateFlag:   proto.Uint32(1),
		CreateParent: proto.Bool(false),
		Replication:  proto.Uint32(defaults.GetReplication()),
		BlockSize:    proto.Uint64(defaults.GetBlockSize()),
	}
	createResp := &hdfs.CreateResponseProto{}

	err = c.namenode.Execute("create", createReq, createResp)
	if err != nil {
		if nnErr, ok := err.(*rpc.NamenodeError); ok {
			err = interpretException(nnErr.Exception, err)
		}

		return &os.PathError{"create", filename, err}
	}

	completeReq := &hdfs.CompleteRequestProto{
		Src:        proto.String(filename),
		ClientName: proto.String(rpc.ClientName),
	}
	completeResp := &hdfs.CompleteResponseProto{}

	err = c.namenode.Execute("complete", completeReq, completeResp)
	if err != nil {
		return &os.PathError{"create", filename, err}
	}

	return nil
}

func (c *Client) fetchDefaults() (*hdfs.FsServerDefaultsProto, error) {
	if c.defaults != nil {
		return c.defaults, nil
	}

	req := &hdfs.GetServerDefaultsRequestProto{}
	resp := &hdfs.GetServerDefaultsResponseProto{}

	err := c.namenode.Execute("getServerDefaults", req, resp)
	if err != nil {
		return nil, err
	}

	c.defaults = resp.GetServerDefaults()
	return c.defaults, nil
}
<|endoftext|>


<commit_before>//go:generate go run xenapi.go

package xenAPI

import (
	"crypto/tls"
	"fmt"
	"github.com/amfranz/go-xmlrpc-client"
	"net/http"
)

type APIResult struct {
	Value interface{}
}

func (client *Client) APICall(method string, params ...interface{}) (result APIResult, err error) {
	rpcParams := xmlrpc.Params{
		Params: params,
	}

	rpcResult := xmlrpc.Struct{}

	err = client.rpc.Call(method, rpcParams, &rpcResult)
	if err != nil {
		return
	}

	status, ok := rpcResult["Status"].(string)
	if !ok {
		err = fmt.Errorf("Expected a field named %q with a string value in the response", "Status")
		return
	}

	if status != "Success" {
		details := rpcResult["ErrorDescription"].([]interface{})
		_objtype := ""
		if details[1] != nil {
			_objtype = details[1].(string)
		}
		var _uuid string
		if len(details) > 2 {
			_uuid = details[2].(string)
		}
		err = &Error{
			code:    details[0].(string),
			objtype: _objtype, // might be nil
			uuid:    _uuid, // optional
		}
		return
	}

	result.Value = rpcResult["Value"]
	return
}

func NewClient(url string, transport *http.Transport) (*Client, error) {
	if transport == nil {
		transport = &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
	}

	rpc, err := xmlrpc.NewClient(url, transport)
	if err != nil {
		return nil, err
	}

	return prepClient(rpc), nil
}
<commit_msg>check number of parameters<commit_after>//go:generate go run xenapi.go

package xenAPI

import (
	"crypto/tls"
	"fmt"
	"github.com/amfranz/go-xmlrpc-client"
	"net/http"
)

type APIResult struct {
	Value interface{}
}

func (client *Client) APICall(method string, params ...interface{}) (result APIResult, err error) {
	rpcParams := xmlrpc.Params{
		Params: params,
	}

	rpcResult := xmlrpc.Struct{}

	err = client.rpc.Call(method, rpcParams, &rpcResult)
	if err != nil {
		return
	}

	status, ok := rpcResult["Status"].(string)
	if !ok {
		err = fmt.Errorf("Expected a field named %q with a string value in the response", "Status")
		return
	}

	if status != "Success" {
		details := rpcResult["ErrorDescription"].([]interface{})
		_objtype := ""
		if len(details) > 1 && details[1] != nil {
			_objtype = details[1].(string)
		}
		var _uuid string
		if len(details) > 2 {
			_uuid = details[2].(string)
		}
		err = &Error{
			code:    details[0].(string),
			objtype: _objtype, // might be nil
			uuid:    _uuid, // optional
		}
		return
	}

	result.Value = rpcResult["Value"]
	return
}

func NewClient(url string, transport *http.Transport) (*Client, error) {
	if transport == nil {
		transport = &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
	}

	rpc, err := xmlrpc.NewClient(url, transport)
	if err != nil {
		return nil, err
	}

	return prepClient(rpc), nil
}
<|endoftext|>


<commit_before>// Package musicmanager implements a client for managing Google Play
// Music libraries.
package musicmanager

import (
	"fmt"
	"net/http"

	convert "github.com/lxr/go.google.musicmanager/internal/convert"
	mmdspb "github.com/lxr/go.google.musicmanager/internal/download_proto/service"
	mmldpb "github.com/lxr/go.google.musicmanager/internal/locker_proto/data"
	mmssjs "github.com/lxr/go.google.musicmanager/internal/session_json"
	mmudpb "github.com/lxr/go.google.musicmanager/internal/upload_proto/data"
	mmuspb "github.com/lxr/go.google.musicmanager/internal/upload_proto/service"
)

// Scope is the OAuth scope used by this API.
const Scope = "https://www.googleapis.com/auth/musicmanager"

// Client is a Music Manager client.
type Client struct {
	client *http.Client
	id     string
}

// NewClient creates a new Music Manager client with the given device ID
// and underlying HTTP client.  The supplied client must be capable of
// authenticating requests with the OAuth scope Scope.
func NewClient(client *http.Client, deviceID string) (*Client, error) {
	if client == nil {
		return nil, fmt.Errorf("client is nil")
	}
	if deviceID == "" {
		return nil, fmt.Errorf("device ID is empty")
	}
	return &Client{client, deviceID}, nil
}

// Register registers the client as a device with the given name under
// the user's Play Music library.  A client must be registered before
// the other method calls can succeed.  Re-registering a client can be
// used to change its name.  Note that there are limits to how many
// devices one account can have registered at a time, with how many
// accounts one device can be registered, and how many devices a user
// can deregister in a year, so be judicious in using this call.
func (c *Client) Register(name string) error {
	res, err := c.upAuth(&mmuspb.UpAuthRequest{
		UploaderId:   c.id,
		FriendlyName: name,
	})
	if err != nil {
		return err
	}
	if res != mmuspb.UploadResponse_OK {
		return RegisterError(res)
	}
	return nil
}

// ExportTrack returns a short-lived download URL for the given track,
// identified by its server ID.  Downloading the track from this URL
// requires no authentication.
func (c *Client) ExportTrack(id string) (string, error) {
	res, err := c.getDownloadSession(&mmssjs.GetDownloadSessionRequest{
		XDeviceID: c.id,
		SongID:    id,
	})
	if err != nil {
		return "", err
	}
	return res.URL, nil
}

// ListTracks lists the user's tracks from least to most recently
// accessed.  It can optionally filter for purchased or promotional
// tracks only, or for tracks last modified after a given Unix timestamp
// (microsecond precision).  Long responses may be returned in chunks,
// in which case the PageToken field of the TrackList object should
// be given to a new ListTracks call.
func (c *Client) ListTracks(purchasedOnly bool, updatedMin int64, pageToken string) (*TrackList, error) {
	var exportType mmdspb.GetTracksToExportRequest_TracksToExportType
	switch purchasedOnly {
	case true:
		exportType = mmdspb.GetTracksToExportRequest_PURCHASED_AND_PROMOTIONAL
	case false:
		exportType = mmdspb.GetTracksToExportRequest_ALL
	}
	res, err := c.getTracksToExport(&mmdspb.GetTracksToExportRequest{
		ClientId:          c.id,
		ExportType:        exportType,
		UpdatedMin:        updatedMin,
		ContinuationToken: pageToken,
	})
	switch err := err.(type) {
	case nil:
		if res.Status != mmdspb.GetTracksToExportResponse_OK {
			return nil, ListError(res.Status)
		}
	case *RequestError:
		// The Google Play servers respond with 304 Not Modified
		// if no tracks have been modified after the updatedMin
		// timestamp.  This is not exactly an error condition,
		// so we break out of the switch rather than return the
		// error; the remaining function body then arranges to
		// return an empty TrackList.
		if err.Code == http.StatusNotModified {
			break
		}
		return nil, err
	default:
		return nil, err
	}
	trackList := new(TrackList)
	convert.Convert(trackList, res)
	trackList.PurchasedOnly = purchasedOnly
	return trackList, nil
}

// ImportTracks returns short-lived upload URLs for the given tracks.
// MP3 audio data can be PUT to these URLs without authentication.
// Individual tracks can fail, in which case errs[i] contains the
// reason why importing tracks[i] failed.
//
// The Title field of a track to be imported cannot be empty.  The
// server also uses the ClientId field to identify which tracks have
// already been uploaded, so trying to import a track with a client ID
// that has already been uploaded will fail.  Leaving the ClientId field
// empty appears to bypass this server-side check; however, the
// implementation of ImportTracks cannot handle two tracks in a batch
// with the same client ID, even if they are both empty.
func (c *Client) ImportTracks(tracks []*Track) (urls []string, errs []error) {
	// Construct and the client-ID-to-track-index mapping and the
	// initial metadata upload.
	cidm := make(map[string]int)
	trks := make([]*mmldpb.Track, len(tracks))
	errs = make([]error, len(tracks))
	for i, track := range tracks {
		if _, ok := cidm[track.ClientId]; ok {
			errs[i] = fmt.Errorf("trying to import two tracks with the same client-side ID")
			continue
		}
		trks[i] = new(mmldpb.Track)
		convert.Convert(trks[i], track)
		cidm[track.ClientId] = i
	}
	// Upload track metadata.
	res, err := c.uploadMetadata(&mmuspb.UploadMetadataRequest{
		UploaderId: c.id,
		Track:      trks,
	})
	if err != nil {
		for i := range errs {
			if errs[i] == nil {
				errs[i] = err
			}
		}
		return nil, errs
	}
	// Satisfy any requests for track samples and append the
	// responses to tres.
	tres := res.TrackSampleResponse
	if n := len(res.SignedChallengeInfo); n > 0 {
		spls := make([]*mmudpb.TrackSample, n)
		for i, sci := range res.SignedChallengeInfo {
			ci := sci.ChallengeInfo
			j := cidm[ci.ClientTrackId]
			sample := []byte{}
			if sampler := tracks[j].Sampler; sampler != nil {
				sample = sampler(
					int(ci.StartMillis),
					int(ci.DurationMillis),
				)
			}
			spls[i] = &mmudpb.TrackSample{
				Track:               trks[j],
				Sample:              sample,
				SignedChallengeInfo: sci,
				SampleFormat:        mmldpb.Track_MP3,
			}
		}
		sres, err := c.uploadSample(&mmuspb.UploadSampleRequest{
			UploaderId:  c.id,
			TrackSample: spls,
		})
		if err != nil {
			for _, spl := range spls {
				errs[cidm[spl.Track.ClientId]] = err
			}
		} else {
			tres = append(tres, sres.TrackSampleResponse...)
		}
	}
	// Parse responses to track metadata and samples.  The result
	// is a map from track indices to their server IDs.
	sidm := make(map[int]string)
	for _, res := range tres {
		i := cidm[res.ClientTrackId]
		if res.ResponseCode != mmuspb.TrackSampleResponse_UPLOAD_REQUESTED {
			errs[i] = ImportError(res.ResponseCode)
			continue
		}
		sidm[i] = res.ServerTrackId
	}
	// Acquire upload sessions.
	uploaded := 0
	total := len(sidm)
	urls = make([]string, len(tracks))
	for i, id := range sidm {
		trk := tracks[i]
		res, err := c.getUploadSession(&mmssjs.GetUploadSessionRequest{
			Name:                      id,
			UploaderId:                c.id,
			ClientId:                  trk.ClientId,
			ServerId:                  id,
			TrackBitRate:              trk.BitRate,
			CurrentUploadingTrack:     trk.Title,
			CurrentTotalUploadedCount: uploaded,
			ClientTotalSongCount:      total,
			SyncNow:                   true,
		})
		if err != nil {
			errs[i] = err
			continue
		}
		if res.Error != nil {
			errs[i] = res.Error
			continue
		}
		urls[i] = res.Transfers[0].PutUrl
	}
	return urls, errs
}
<commit_msg>Make it safe to return a nil byte slice in Client.Sampler<commit_after>// Package musicmanager implements a client for managing Google Play
// Music libraries.
package musicmanager

import (
	"fmt"
	"net/http"

	convert "github.com/lxr/go.google.musicmanager/internal/convert"
	mmdspb "github.com/lxr/go.google.musicmanager/internal/download_proto/service"
	mmldpb "github.com/lxr/go.google.musicmanager/internal/locker_proto/data"
	mmssjs "github.com/lxr/go.google.musicmanager/internal/session_json"
	mmudpb "github.com/lxr/go.google.musicmanager/internal/upload_proto/data"
	mmuspb "github.com/lxr/go.google.musicmanager/internal/upload_proto/service"
)

// Scope is the OAuth scope used by this API.
const Scope = "https://www.googleapis.com/auth/musicmanager"

// Client is a Music Manager client.
type Client struct {
	client *http.Client
	id     string
}

// NewClient creates a new Music Manager client with the given device ID
// and underlying HTTP client.  The supplied client must be capable of
// authenticating requests with the OAuth scope Scope.
func NewClient(client *http.Client, deviceID string) (*Client, error) {
	if client == nil {
		return nil, fmt.Errorf("client is nil")
	}
	if deviceID == "" {
		return nil, fmt.Errorf("device ID is empty")
	}
	return &Client{client, deviceID}, nil
}

// Register registers the client as a device with the given name under
// the user's Play Music library.  A client must be registered before
// the other method calls can succeed.  Re-registering a client can be
// used to change its name.  Note that there are limits to how many
// devices one account can have registered at a time, with how many
// accounts one device can be registered, and how many devices a user
// can deregister in a year, so be judicious in using this call.
func (c *Client) Register(name string) error {
	res, err := c.upAuth(&mmuspb.UpAuthRequest{
		UploaderId:   c.id,
		FriendlyName: name,
	})
	if err != nil {
		return err
	}
	if res != mmuspb.UploadResponse_OK {
		return RegisterError(res)
	}
	return nil
}

// ExportTrack returns a short-lived download URL for the given track,
// identified by its server ID.  Downloading the track from this URL
// requires no authentication.
func (c *Client) ExportTrack(id string) (string, error) {
	res, err := c.getDownloadSession(&mmssjs.GetDownloadSessionRequest{
		XDeviceID: c.id,
		SongID:    id,
	})
	if err != nil {
		return "", err
	}
	return res.URL, nil
}

// ListTracks lists the user's tracks from least to most recently
// accessed.  It can optionally filter for purchased or promotional
// tracks only, or for tracks last modified after a given Unix timestamp
// (microsecond precision).  Long responses may be returned in chunks,
// in which case the PageToken field of the TrackList object should
// be given to a new ListTracks call.
func (c *Client) ListTracks(purchasedOnly bool, updatedMin int64, pageToken string) (*TrackList, error) {
	var exportType mmdspb.GetTracksToExportRequest_TracksToExportType
	switch purchasedOnly {
	case true:
		exportType = mmdspb.GetTracksToExportRequest_PURCHASED_AND_PROMOTIONAL
	case false:
		exportType = mmdspb.GetTracksToExportRequest_ALL
	}
	res, err := c.getTracksToExport(&mmdspb.GetTracksToExportRequest{
		ClientId:          c.id,
		ExportType:        exportType,
		UpdatedMin:        updatedMin,
		ContinuationToken: pageToken,
	})
	switch err := err.(type) {
	case nil:
		if res.Status != mmdspb.GetTracksToExportResponse_OK {
			return nil, ListError(res.Status)
		}
	case *RequestError:
		// The Google Play servers respond with 304 Not Modified
		// if no tracks have been modified after the updatedMin
		// timestamp.  This is not exactly an error condition,
		// so we break out of the switch rather than return the
		// error; the remaining function body then arranges to
		// return an empty TrackList.
		if err.Code == http.StatusNotModified {
			break
		}
		return nil, err
	default:
		return nil, err
	}
	trackList := new(TrackList)
	convert.Convert(trackList, res)
	trackList.PurchasedOnly = purchasedOnly
	return trackList, nil
}

// ImportTracks returns short-lived upload URLs for the given tracks.
// MP3 audio data can be PUT to these URLs without authentication.
// Individual tracks can fail, in which case errs[i] contains the
// reason why importing tracks[i] failed.
//
// The Title field of a track to be imported cannot be empty.  The
// server also uses the ClientId field to identify which tracks have
// already been uploaded, so trying to import a track with a client ID
// that has already been uploaded will fail.  Leaving the ClientId field
// empty appears to bypass this server-side check; however, the
// implementation of ImportTracks cannot handle two tracks in a batch
// with the same client ID, even if they are both empty.
func (c *Client) ImportTracks(tracks []*Track) (urls []string, errs []error) {
	// Construct and the client-ID-to-track-index mapping and the
	// initial metadata upload.
	cidm := make(map[string]int)
	trks := make([]*mmldpb.Track, len(tracks))
	errs = make([]error, len(tracks))
	for i, track := range tracks {
		if _, ok := cidm[track.ClientId]; ok {
			errs[i] = fmt.Errorf("trying to import two tracks with the same client-side ID")
			continue
		}
		trks[i] = new(mmldpb.Track)
		convert.Convert(trks[i], track)
		cidm[track.ClientId] = i
	}
	// Upload track metadata.
	res, err := c.uploadMetadata(&mmuspb.UploadMetadataRequest{
		UploaderId: c.id,
		Track:      trks,
	})
	if err != nil {
		for i := range errs {
			if errs[i] == nil {
				errs[i] = err
			}
		}
		return nil, errs
	}
	// Satisfy any requests for track samples and append the
	// responses to tres.
	tres := res.TrackSampleResponse
	if n := len(res.SignedChallengeInfo); n > 0 {
		spls := make([]*mmudpb.TrackSample, n)
		for i, sci := range res.SignedChallengeInfo {
			ci := sci.ChallengeInfo
			j := cidm[ci.ClientTrackId]
			var sample []byte
			if sampler := tracks[j].Sampler; sampler != nil {
				sample = sampler(
					int(ci.StartMillis),
					int(ci.DurationMillis),
				)
			}
			if sample == nil {
				// A nil sample is different from an
				// empty one: the former results in an
				// invalid sample message.  So, if
				// Sampler leaves us with a nil sample,
				// replace it with an empty one.
				sample = make([]byte, 0)
			}
			spls[i] = &mmudpb.TrackSample{
				Track:               trks[j],
				Sample:              sample,
				SignedChallengeInfo: sci,
				SampleFormat:        mmldpb.Track_MP3,
			}
		}
		sres, err := c.uploadSample(&mmuspb.UploadSampleRequest{
			UploaderId:  c.id,
			TrackSample: spls,
		})
		if err != nil {
			for _, spl := range spls {
				errs[cidm[spl.Track.ClientId]] = err
			}
		} else {
			tres = append(tres, sres.TrackSampleResponse...)
		}
	}
	// Parse responses to track metadata and samples.  The result
	// is a map from track indices to their server IDs.
	sidm := make(map[int]string)
	for _, res := range tres {
		i := cidm[res.ClientTrackId]
		if res.ResponseCode != mmuspb.TrackSampleResponse_UPLOAD_REQUESTED {
			errs[i] = ImportError(res.ResponseCode)
			continue
		}
		sidm[i] = res.ServerTrackId
	}
	// Acquire upload sessions.
	uploaded := 0
	total := len(sidm)
	urls = make([]string, len(tracks))
	for i, id := range sidm {
		trk := tracks[i]
		res, err := c.getUploadSession(&mmssjs.GetUploadSessionRequest{
			Name:                      id,
			UploaderId:                c.id,
			ClientId:                  trk.ClientId,
			ServerId:                  id,
			TrackBitRate:              trk.BitRate,
			CurrentUploadingTrack:     trk.Title,
			CurrentTotalUploadedCount: uploaded,
			ClientTotalSongCount:      total,
			SyncNow:                   true,
		})
		if err != nil {
			errs[i] = err
			continue
		}
		if res.Error != nil {
			errs[i] = res.Error
			continue
		}
		urls[i] = res.Transfers[0].PutUrl
	}
	return urls, errs
}
<|endoftext|>


<commit_before>package greq

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/base64"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"

	requestwork "github.com/syhlion/requestwork.v2"
)

type Trace struct {
	Url              string        `json:"url"`
	Method           string        `json:"method"`
	Body             string        `json:"body"`
	Param            string        `json:"param"`
	DNSLookup        time.Duration `json:"dns_lookup"`
	TCPConnection    time.Duration `json:"tcp_connection"`
	TLSHandshake     time.Duration `json:"tls_handshake"`
	ServerProcessing time.Duration `json:"server_prcoessing"`
	ContentTransfer  time.Duration `json:"content_transfer"`
	NameLookup       time.Duration `json:"name_lookup"`
	Connect          time.Duration `json:"connect"`
	PreTransfer      time.Duration `json:"pre_transfer"`
	StartTransfer    time.Duration `json:"start_transfer"`
	Total            time.Duration `json:"total"`
}

var ip string

func getExternalIP() (string, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", err
	}
	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}
		addrs, err := iface.Addrs()
		if err != nil {
			return "", err
		}
		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}
			if ip == nil || ip.IsLoopback() {
				continue
			}
			ip = ip.To4()
			if ip == nil {
				continue // not an ipv4 address
			}
			return ip.String(), nil
		}
	}
	return "", errors.New("are you connected to the network?")
}

func init() {
	log.SetFormatter(&log.JSONFormatter{})

	log.SetOutput(os.Stdout)
	log.SetLevel(log.DebugLevel)
	ip, _ = getExternalIP()
}

//New return http client
func New(worker *requestwork.Worker, timeout time.Duration, debug bool) *Client {
	return &Client{
		worker:  worker,
		timeout: timeout,
		headers: make(map[string]string),
		lock:    &sync.RWMutex{},
		debug:   debug,
	}
}

//Client instance
type Client struct {
	worker  *requestwork.Worker
	timeout time.Duration
	headers map[string]string
	host    string
	lock    *sync.RWMutex
	debug   bool
}

func (c *Client) CheckRedircet(f func(req *http.Request, via []*http.Request) error) {
	c.worker.CheckRedirect(f)
}

//SetBasicAuth  set Basic auth
func (c *Client) SetBasicAuth(username, password string) *Client {
	auth := username + ":" + password
	hash := base64.StdEncoding.EncodeToString([]byte(auth))
	c.lock.Lock()
	defer c.lock.Unlock()
	c.headers["Authorization"] = "Basic " + hash
	return c
}

//SetHeader set http header
func (c *Client) SetHeader(key, value string) *Client {
	key = strings.Title(key)
	c.lock.Lock()
	defer c.lock.Unlock()
	c.headers[key] = value
	return c
}

//SetHost set host
func (c *Client) SetHost(host string) *Client {
	c.lock.Lock()
	defer c.lock.Unlock()

	c.host = host
	return c
}

//Get http method get
func (c *Client) Get(url string, params url.Values) (data []byte, httpstatus int, err error) {
	if params != nil {
		url += "?" + params.Encode()
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	return c.resolveRequest(req, params, err)

}
func (c *Client) GetWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	if params != nil {
		url += "?" + params.Encode()
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRequest(req, params, err)
}

//Post http method post
func (c *Client) Post(url string, params url.Values) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(params.Encode()))
	return c.resolveRequest(req, params, err)
}
func (c *Client) PostRaw(url string, body io.Reader) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, body)
	return c.resolveRawRequest(req, body, err)
}
func (c *Client) PostRawWithOnceHeader(url string, body io.Reader, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, body)
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRawRequest(req, body, err)
}

func (c *Client) PostWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(params.Encode()))
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRequest(req, params, err)
}

//Put http method put
func (c *Client) Put(url string, params url.Values) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPut, url, strings.NewReader(params.Encode()))
	return c.resolveRequest(req, params, err)
}
func (c *Client) PutWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPut, url, strings.NewReader(params.Encode()))
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}

	return c.resolveRequest(req, params, err)
}

//Delete http method Delete
func (c *Client) Delete(url string, params url.Values) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodDelete, url, strings.NewReader(params.Encode()))
	return c.resolveRequest(req, params, err)
}
func (c *Client) DeleteWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(params.Encode()))
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRequest(req, params, err)
}

func (c *Client) resolveHeaders(req *http.Request) {
	c.lock.RLock()
	defer c.lock.RUnlock()
	for key, value := range c.headers {
		req.Header.Set(key, value)
	}
	if c.host != "" {
		req.Host = c.host
	}
}

func (c *Client) resolveRawRequest(req *http.Request, bb io.Reader, e error) (data []byte, httpstatus int, err error) {
	var (
		body           []byte
		status         int
		trace          *httptrace.ClientTrace
		t0, t3, t4, t5 time.Time
	)
	req.Close = true
	t0 = time.Now()
	if c.debug {
		var stat Trace
		defer func() {
			buf := new(bytes.Buffer)
			buf.ReadFrom(bb)
			stat.Param = buf.String()
			stat.Url = req.URL.String()
			stat.Method = req.Method
			stat.Body = string(data)
			stat.TCPConnection = t3.Sub(t0)
			stat.ServerProcessing = t4.Sub(t3)
			stat.ContentTransfer = t5.Sub(t4)
			stat.Connect = t3.Sub(t0)
			stat.StartTransfer = t4.Sub(t0)
			stat.Total = t5.Sub(t0)
			log.WithFields(log.Fields{
				"ip":     ip,
				"name":   "syhlion/greq",
				"param":  stat.Param,
				"url":    stat.Url,
				"method": stat.Method,
				//	"body":              stat.Body,
				"tcp_connection":    stat.TCPConnection.String(),
				"server_processing": stat.ServerProcessing.String(),
				"content_transfer":  stat.ContentTransfer.String(),
				"connect":           stat.Connect.String(),
				"start_transfer":    stat.StartTransfer.String(),
				"total":             stat.Total.String(),
			}).Debug("http trace")

		}()
		trace = &httptrace.ClientTrace{
			GotConn:              func(_ httptrace.GotConnInfo) { t3 = time.Now() },
			GotFirstResponseByte: func() { t4 = time.Now() },
		}
		req = req.WithContext(httptrace.WithClientTrace(context.Background(), trace))
	}
	if e != nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)

	defer cancel()
	c.resolveHeaders(req)
	switch req.Method {
	case "PUT", "POST", "DELETE":
		tmp := req.Header.Get("Content-Type")
		if tmp == "" {
			req.Header.Set("Content-Type", "application/json")
		}
	}
	req = req.WithContext(ctx)

	err = c.worker.Execute(req, func(resp *http.Response, err error) (er error) {
		if err != nil {
			return err
		}
		var readErr error
		defer func() {
			resp.Body.Close()
			t5 = time.Now()
		}()
		status = resp.StatusCode
		switch resp.Header.Get("Content-Encoding") {
		case "gzip":
			reader, err := gzip.NewReader(resp.Body)
			if err != nil {
				return err
			}
			defer reader.Close()
			body, readErr = ioutil.ReadAll(reader)
			if readErr != nil {
				return readErr
			}
		default:
			body, readErr = ioutil.ReadAll(resp.Body)
			if readErr != nil {
				return readErr
			}
		}
		return
	})
	if err != nil {
		return
	}
	data = body
	httpstatus = status
	return

}
func (c *Client) resolveRequest(req *http.Request, params url.Values, e error) (data []byte, httpstatus int, err error) {
	var (
		body           []byte
		status         int
		trace          *httptrace.ClientTrace
		t0, t3, t4, t5 time.Time
	)
	req.Close = true
	t0 = time.Now()
	if c.debug {
		var stat Trace
		defer func() {
			stat.Param = params.Encode()
			stat.Url = req.URL.String()
			stat.Method = req.Method
			stat.Body = string(data)
			stat.TCPConnection = t3.Sub(t0)
			stat.ServerProcessing = t4.Sub(t3)
			stat.ContentTransfer = t5.Sub(t4)
			stat.Connect = t3.Sub(t0)
			stat.StartTransfer = t4.Sub(t0)
			stat.Total = t5.Sub(t0)
			log.WithFields(log.Fields{
				"ip":     ip,
				"name":   "syhlion/greq",
				"param":  stat.Param,
				"url":    stat.Url,
				"method": stat.Method,
				//	"body":              stat.Body,
				"tcp_connection":    stat.TCPConnection.String(),
				"server_processing": stat.ServerProcessing.String(),
				"content_transfer":  stat.ContentTransfer.String(),
				"connect":           stat.Connect.String(),
				"start_transfer":    stat.StartTransfer.String(),
				"total":             stat.Total.String(),
			}).Debug("http trace")

		}()
		trace = &httptrace.ClientTrace{
			GotConn:              func(_ httptrace.GotConnInfo) { t3 = time.Now() },
			GotFirstResponseByte: func() { t4 = time.Now() },
		}
		req = req.WithContext(httptrace.WithClientTrace(context.Background(), trace))
	}
	if e != nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)

	defer cancel()
	c.resolveHeaders(req)

	switch req.Method {
	case "PUT", "POST", "DELETE":
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded; charset=UTF-8")
	}
	req = req.WithContext(ctx)

	err = c.worker.Execute(req, func(resp *http.Response, err error) (er error) {
		if err != nil {
			return err
		}
		var readErr error
		defer func() {
			resp.Body.Close()
			t5 = time.Now()
		}()
		status = resp.StatusCode
		switch resp.Header.Get("Content-Encoding") {
		case "gzip":
			reader, err := gzip.NewReader(resp.Body)
			if err != nil {
				return err
			}
			defer reader.Close()
			body, readErr = ioutil.ReadAll(reader)
			if readErr != nil {
				return readErr
			}
		default:
			body, readErr = ioutil.ReadAll(resp.Body)
			if readErr != nil {
				return readErr
			}
		}
		return
	})
	if err != nil {
		return
	}
	data = body
	httpstatus = status
	return

}

func (c *Client) ResolveRequest(req *http.Request, params url.Values, e error) (data []byte, httpstatus int, err error) {
	return c.resolveRequest(req, params, err)
}
<commit_msg>add putraw<commit_after>package greq

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/base64"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"

	requestwork "github.com/syhlion/requestwork.v2"
)

type Trace struct {
	Url              string        `json:"url"`
	Method           string        `json:"method"`
	Body             string        `json:"body"`
	Param            string        `json:"param"`
	DNSLookup        time.Duration `json:"dns_lookup"`
	TCPConnection    time.Duration `json:"tcp_connection"`
	TLSHandshake     time.Duration `json:"tls_handshake"`
	ServerProcessing time.Duration `json:"server_prcoessing"`
	ContentTransfer  time.Duration `json:"content_transfer"`
	NameLookup       time.Duration `json:"name_lookup"`
	Connect          time.Duration `json:"connect"`
	PreTransfer      time.Duration `json:"pre_transfer"`
	StartTransfer    time.Duration `json:"start_transfer"`
	Total            time.Duration `json:"total"`
}

var ip string

func getExternalIP() (string, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", err
	}
	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // loopback interface
		}
		addrs, err := iface.Addrs()
		if err != nil {
			return "", err
		}
		for _, addr := range addrs {
			var ip net.IP
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
			case *net.IPAddr:
				ip = v.IP
			}
			if ip == nil || ip.IsLoopback() {
				continue
			}
			ip = ip.To4()
			if ip == nil {
				continue // not an ipv4 address
			}
			return ip.String(), nil
		}
	}
	return "", errors.New("are you connected to the network?")
}

func init() {
	log.SetFormatter(&log.JSONFormatter{})

	log.SetOutput(os.Stdout)
	log.SetLevel(log.DebugLevel)
	ip, _ = getExternalIP()
}

//New return http client
func New(worker *requestwork.Worker, timeout time.Duration, debug bool) *Client {
	return &Client{
		worker:  worker,
		timeout: timeout,
		headers: make(map[string]string),
		lock:    &sync.RWMutex{},
		debug:   debug,
	}
}

//Client instance
type Client struct {
	worker  *requestwork.Worker
	timeout time.Duration
	headers map[string]string
	host    string
	lock    *sync.RWMutex
	debug   bool
}

func (c *Client) CheckRedircet(f func(req *http.Request, via []*http.Request) error) {
	c.worker.CheckRedirect(f)
}

//SetBasicAuth  set Basic auth
func (c *Client) SetBasicAuth(username, password string) *Client {
	auth := username + ":" + password
	hash := base64.StdEncoding.EncodeToString([]byte(auth))
	c.lock.Lock()
	defer c.lock.Unlock()
	c.headers["Authorization"] = "Basic " + hash
	return c
}

//SetHeader set http header
func (c *Client) SetHeader(key, value string) *Client {
	key = strings.Title(key)
	c.lock.Lock()
	defer c.lock.Unlock()
	c.headers[key] = value
	return c
}

//SetHost set host
func (c *Client) SetHost(host string) *Client {
	c.lock.Lock()
	defer c.lock.Unlock()

	c.host = host
	return c
}

//Get http method get
func (c *Client) Get(url string, params url.Values) (data []byte, httpstatus int, err error) {
	if params != nil {
		url += "?" + params.Encode()
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	return c.resolveRequest(req, params, err)

}
func (c *Client) GetWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	if params != nil {
		url += "?" + params.Encode()
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRequest(req, params, err)
}

//Post http method post
func (c *Client) Post(url string, params url.Values) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(params.Encode()))
	return c.resolveRequest(req, params, err)
}
func (c *Client) PostRaw(url string, body io.Reader) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, body)
	return c.resolveRawRequest(req, body, err)
}
func (c *Client) PostRawWithOnceHeader(url string, body io.Reader, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, body)
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRawRequest(req, body, err)
}

func (c *Client) PostWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(params.Encode()))
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRequest(req, params, err)
}

//Put http method put
func (c *Client) Put(url string, params url.Values) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPut, url, strings.NewReader(params.Encode()))
	return c.resolveRequest(req, params, err)
}
func (c *Client) PutRaw(url string, body io.Reader) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPut, url, body)
	return c.resolveRawRequest(req, body, err)
}
func (c *Client) PutWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPut, url, strings.NewReader(params.Encode()))
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}

	return c.resolveRequest(req, params, err)
}

//Delete http method Delete
func (c *Client) Delete(url string, params url.Values) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodDelete, url, strings.NewReader(params.Encode()))
	return c.resolveRequest(req, params, err)
}
func (c *Client) DeleteWithOnceHeader(url string, params url.Values, headers map[string]string) (data []byte, httpstatus int, err error) {
	req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(params.Encode()))
	if err != nil {
		return
	}
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	return c.resolveRequest(req, params, err)
}

func (c *Client) resolveHeaders(req *http.Request) {
	c.lock.RLock()
	defer c.lock.RUnlock()
	for key, value := range c.headers {
		req.Header.Set(key, value)
	}
	if c.host != "" {
		req.Host = c.host
	}
}

func (c *Client) resolveRawRequest(req *http.Request, bb io.Reader, e error) (data []byte, httpstatus int, err error) {
	var (
		body           []byte
		status         int
		trace          *httptrace.ClientTrace
		t0, t3, t4, t5 time.Time
	)
	req.Close = true
	t0 = time.Now()
	if c.debug {
		var stat Trace
		defer func() {
			buf := new(bytes.Buffer)
			buf.ReadFrom(bb)
			stat.Param = buf.String()
			stat.Url = req.URL.String()
			stat.Method = req.Method
			stat.Body = string(data)
			stat.TCPConnection = t3.Sub(t0)
			stat.ServerProcessing = t4.Sub(t3)
			stat.ContentTransfer = t5.Sub(t4)
			stat.Connect = t3.Sub(t0)
			stat.StartTransfer = t4.Sub(t0)
			stat.Total = t5.Sub(t0)
			log.WithFields(log.Fields{
				"ip":     ip,
				"name":   "syhlion/greq",
				"param":  stat.Param,
				"url":    stat.Url,
				"method": stat.Method,
				//	"body":              stat.Body,
				"tcp_connection":    stat.TCPConnection.String(),
				"server_processing": stat.ServerProcessing.String(),
				"content_transfer":  stat.ContentTransfer.String(),
				"connect":           stat.Connect.String(),
				"start_transfer":    stat.StartTransfer.String(),
				"total":             stat.Total.String(),
			}).Debug("http trace")

		}()
		trace = &httptrace.ClientTrace{
			GotConn:              func(_ httptrace.GotConnInfo) { t3 = time.Now() },
			GotFirstResponseByte: func() { t4 = time.Now() },
		}
		req = req.WithContext(httptrace.WithClientTrace(context.Background(), trace))
	}
	if e != nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)

	defer cancel()
	c.resolveHeaders(req)
	switch req.Method {
	case "PUT", "POST", "DELETE":
		tmp := req.Header.Get("Content-Type")
		if tmp == "" {
			req.Header.Set("Content-Type", "application/json")
		}
	}
	req = req.WithContext(ctx)

	err = c.worker.Execute(req, func(resp *http.Response, err error) (er error) {
		if err != nil {
			return err
		}
		var readErr error
		defer func() {
			resp.Body.Close()
			t5 = time.Now()
		}()
		status = resp.StatusCode
		switch resp.Header.Get("Content-Encoding") {
		case "gzip":
			reader, err := gzip.NewReader(resp.Body)
			if err != nil {
				return err
			}
			defer reader.Close()
			body, readErr = ioutil.ReadAll(reader)
			if readErr != nil {
				return readErr
			}
		default:
			body, readErr = ioutil.ReadAll(resp.Body)
			if readErr != nil {
				return readErr
			}
		}
		return
	})
	if err != nil {
		return
	}
	data = body
	httpstatus = status
	return

}
func (c *Client) resolveRequest(req *http.Request, params url.Values, e error) (data []byte, httpstatus int, err error) {
	var (
		body           []byte
		status         int
		trace          *httptrace.ClientTrace
		t0, t3, t4, t5 time.Time
	)
	req.Close = true
	t0 = time.Now()
	if c.debug {
		var stat Trace
		defer func() {
			stat.Param = params.Encode()
			stat.Url = req.URL.String()
			stat.Method = req.Method
			stat.Body = string(data)
			stat.TCPConnection = t3.Sub(t0)
			stat.ServerProcessing = t4.Sub(t3)
			stat.ContentTransfer = t5.Sub(t4)
			stat.Connect = t3.Sub(t0)
			stat.StartTransfer = t4.Sub(t0)
			stat.Total = t5.Sub(t0)
			log.WithFields(log.Fields{
				"ip":     ip,
				"name":   "syhlion/greq",
				"param":  stat.Param,
				"url":    stat.Url,
				"method": stat.Method,
				//	"body":              stat.Body,
				"tcp_connection":    stat.TCPConnection.String(),
				"server_processing": stat.ServerProcessing.String(),
				"content_transfer":  stat.ContentTransfer.String(),
				"connect":           stat.Connect.String(),
				"start_transfer":    stat.StartTransfer.String(),
				"total":             stat.Total.String(),
			}).Debug("http trace")

		}()
		trace = &httptrace.ClientTrace{
			GotConn:              func(_ httptrace.GotConnInfo) { t3 = time.Now() },
			GotFirstResponseByte: func() { t4 = time.Now() },
		}
		req = req.WithContext(httptrace.WithClientTrace(context.Background(), trace))
	}
	if e != nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)

	defer cancel()
	c.resolveHeaders(req)

	switch req.Method {
	case "PUT", "POST", "DELETE":
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded; charset=UTF-8")
	}
	req = req.WithContext(ctx)

	err = c.worker.Execute(req, func(resp *http.Response, err error) (er error) {
		if err != nil {
			return err
		}
		var readErr error
		defer func() {
			resp.Body.Close()
			t5 = time.Now()
		}()
		status = resp.StatusCode
		switch resp.Header.Get("Content-Encoding") {
		case "gzip":
			reader, err := gzip.NewReader(resp.Body)
			if err != nil {
				return err
			}
			defer reader.Close()
			body, readErr = ioutil.ReadAll(reader)
			if readErr != nil {
				return readErr
			}
		default:
			body, readErr = ioutil.ReadAll(resp.Body)
			if readErr != nil {
				return readErr
			}
		}
		return
	})
	if err != nil {
		return
	}
	data = body
	httpstatus = status
	return

}

func (c *Client) ResolveRequest(req *http.Request, params url.Values, e error) (data []byte, httpstatus int, err error) {
	return c.resolveRequest(req, params, err)
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"os"
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	"golang.org/x/net/context"
)

// KeybaseDaemonRPC implements the KeybaseService interface using RPC
// calls.
type KeybaseDaemonRPC struct {
	*KeybaseServiceBase
	daemonLog logger.Logger

	// Only used when there's a real connection (i.e., not in
	// testing).
	shutdownFn func()

	keepAliveCancel context.CancelFunc

	// protocols (additional to required protocols) to register on server connect
	protocols []rpc.Protocol

	// simplefs is the simplefs implementation used (if not nil)
	simplefs keybase1.SimpleFSInterface
}

var _ keybase1.NotifySessionInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyKeyfamilyInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyPaperKeyInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyFSRequestInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyTeamInterface = (*KeybaseDaemonRPC)(nil)

var _ rpc.ConnectionHandler = (*KeybaseDaemonRPC)(nil)

var _ KeybaseService = (*KeybaseDaemonRPC)(nil)

// NewKeybaseDaemonRPC makes a new KeybaseDaemonRPC that makes RPC
// calls using the socket of the given Keybase context.
func NewKeybaseDaemonRPC(config Config, kbCtx Context, log logger.Logger,
	debug bool, createSimpleFS func(Config) keybase1.SimpleFSInterface,
) *KeybaseDaemonRPC {
	k := newKeybaseDaemonRPC(config, kbCtx, log)
	k.config = config
	k.daemonLog = logger.NewWithCallDepth("daemon", 1)
	if createSimpleFS != nil {
		k.simplefs = createSimpleFS(config)
	}
	if debug {
		k.daemonLog.Configure("", true, "")
	}
	conn := NewSharedKeybaseConnection(kbCtx, config, k)
	k.fillClients(conn.GetClient())
	k.shutdownFn = conn.Shutdown

	if config.Mode() != InitMinimal {
		ctx, cancel := context.WithCancel(context.Background())
		k.keepAliveCancel = cancel
		go k.keepAliveLoop(ctx)
	}

	return k
}

// For testing.
func newKeybaseDaemonRPCWithClient(kbCtx Context, client rpc.GenericClient,
	log logger.Logger) *KeybaseDaemonRPC {
	k := newKeybaseDaemonRPC(nil, kbCtx, log)
	k.fillClients(client)
	// No need for a keepalive loop in this case, since this is only
	// used during testing.
	return k
}

func newKeybaseDaemonRPC(config Config, kbCtx Context, log logger.Logger) *KeybaseDaemonRPC {
	serviceBase := NewKeybaseServiceBase(config, kbCtx, log)
	if serviceBase == nil {
		return nil
	}
	k := KeybaseDaemonRPC{
		KeybaseServiceBase: serviceBase,
	}
	return &k
}

func (k *KeybaseDaemonRPC) fillClients(client rpc.GenericClient) {
	k.FillClients(keybase1.IdentifyClient{Cli: client},
		keybase1.UserClient{Cli: client},
		keybase1.TeamsClient{Cli: client},
		keybase1.MerkleClient{Cli: client},
		keybase1.SessionClient{Cli: client},
		keybase1.FavoriteClient{Cli: client},
		keybase1.KbfsClient{Cli: client},
		keybase1.KbfsMountClient{Cli: client})
}

type daemonLogUI struct {
	log logger.Logger
}

var _ keybase1.LogUiInterface = daemonLogUI{}

func (l daemonLogUI) Log(ctx context.Context, arg keybase1.LogArg) error {
	format := "%s"
	// arg.Text.Markup should always be false, so ignore it.
	s := arg.Text.Data
	switch arg.Level {
	case keybase1.LogLevel_DEBUG:
		l.log.CDebugf(ctx, format, s)
	case keybase1.LogLevel_INFO:
		l.log.CInfof(ctx, format, s)
	case keybase1.LogLevel_WARN:
		l.log.CWarningf(ctx, format, s)
	case keybase1.LogLevel_ERROR:
		l.log.CErrorf(ctx, format, s)
	case keybase1.LogLevel_NOTICE:
		l.log.CNoticef(ctx, format, s)
	case keybase1.LogLevel_CRITICAL:
		l.log.CCriticalf(ctx, format, s)
	default:
		l.log.CWarningf(ctx, format, s)
	}
	return nil
}

type daemonIdentifyUI struct {
	log logger.Logger
}

var _ keybase1.IdentifyUiInterface = daemonIdentifyUI{}

func (d daemonIdentifyUI) DelegateIdentifyUI(ctx context.Context) (int, error) {
	d.log.CDebugf(ctx, "DelegateIdentifyUI() (returning 0, UIDelegationUnavailableError{}")
	return 0, libkb.UIDelegationUnavailableError{}
}

func (d daemonIdentifyUI) Start(ctx context.Context, arg keybase1.StartArg) error {
	d.log.CDebugf(ctx, "Start(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayKey(ctx context.Context, arg keybase1.DisplayKeyArg) error {
	d.log.CDebugf(ctx, "DisplayKey(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) ReportLastTrack(ctx context.Context, arg keybase1.ReportLastTrackArg) error {
	d.log.CDebugf(ctx, "ReportLastTrack(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) LaunchNetworkChecks(ctx context.Context, arg keybase1.LaunchNetworkChecksArg) error {
	d.log.CDebugf(ctx, "LaunchNetworkChecks(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayTrackStatement(ctx context.Context, arg keybase1.DisplayTrackStatementArg) error {
	d.log.CDebugf(ctx, "DisplayTrackStatement(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) FinishWebProofCheck(ctx context.Context, arg keybase1.FinishWebProofCheckArg) error {
	d.log.CDebugf(ctx, "FinishWebProofCheck(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) FinishSocialProofCheck(ctx context.Context, arg keybase1.FinishSocialProofCheckArg) error {
	d.log.CDebugf(ctx, "FinishSocialProofCheck(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayCryptocurrency(ctx context.Context, arg keybase1.DisplayCryptocurrencyArg) error {
	d.log.CDebugf(ctx, "DisplayCryptocurrency(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) ReportTrackToken(ctx context.Context, arg keybase1.ReportTrackTokenArg) error {
	d.log.CDebugf(ctx, "ReportTrackToken(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayUserCard(ctx context.Context, arg keybase1.DisplayUserCardArg) error {
	d.log.CDebugf(ctx, "DisplayUserCard(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) Confirm(ctx context.Context, arg keybase1.ConfirmArg) (keybase1.ConfirmResult, error) {
	d.log.CDebugf(ctx, "Confirm(%+v) (returning false)", arg)
	return keybase1.ConfirmResult{
		IdentityConfirmed: false,
		RemoteConfirmed:   false,
	}, nil
}

func (d daemonIdentifyUI) DisplayTLFCreateWithInvite(ctx context.Context,
	arg keybase1.DisplayTLFCreateWithInviteArg) (err error) {
	return nil
}

func (d daemonIdentifyUI) Dismiss(ctx context.Context,
	_ keybase1.DismissArg) error {
	return nil
}

func (d daemonIdentifyUI) Cancel(ctx context.Context, sessionID int) error {
	d.log.CDebugf(ctx, "Cancel(%d)", sessionID)
	return nil
}

func (d daemonIdentifyUI) Finish(ctx context.Context, sessionID int) error {
	d.log.CDebugf(ctx, "Finish(%d)", sessionID)
	return nil
}

// HandlerName implements the ConnectionHandler interface.
func (*KeybaseDaemonRPC) HandlerName() string {
	return "KeybaseDaemonRPC"
}

// AddProtocols adds protocols that are registered on server connect
func (k *KeybaseDaemonRPC) AddProtocols(protocols []rpc.Protocol) {
	if protocols == nil {
		return
	}
	if k.protocols != nil {
		k.protocols = append(k.protocols, protocols...)
	} else {
		k.protocols = protocols
	}
}

// OnConnect implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnConnect(ctx context.Context,
	conn *rpc.Connection, rawClient rpc.GenericClient,
	server *rpc.Server) error {

	// Protocols that KBFS requires
	protocols := []rpc.Protocol{
		keybase1.LogUiProtocol(daemonLogUI{k.daemonLog}),
		keybase1.IdentifyUiProtocol(daemonIdentifyUI{k.daemonLog}),
		keybase1.NotifySessionProtocol(k),
		keybase1.NotifyKeyfamilyProtocol(k),
		keybase1.NotifyPaperKeyProtocol(k),
		keybase1.NotifyFSRequestProtocol(k),
		keybase1.NotifyTeamProtocol(k),
		keybase1.TlfKeysProtocol(k),
		keybase1.ReachabilityProtocol(k),
	}

	// Add simplefs if set
	if k.simplefs != nil {
		protocols = append(protocols, keybase1.SimpleFSProtocol(k.simplefs))
	}

	if k.protocols != nil {
		protocols = append(protocols, k.protocols...)
	}

	for _, p := range protocols {
		err := server.Register(p)
		if err != nil {
			if _, ok := err.(rpc.AlreadyRegisteredError); !ok {
				return err
			}
		}
	}

	// Using conn.GetClient() here would cause problematic
	// recursion.
	c := keybase1.NotifyCtlClient{Cli: rawClient}
	err := c.SetNotifications(ctx, keybase1.NotificationChannels{
		Session:      true,
		Paperkeys:    true,
		Keyfamily:    true,
		Kbfsrequest:  true,
		Reachability: true,
		Team:         true,
	})
	if err != nil {
		return err
	}

	// Introduce ourselves. TODO: move this to SharedKeybaseConnection
	// somehow?
	configClient := keybase1.ConfigClient{Cli: rawClient}
	err = configClient.HelloIAm(ctx, keybase1.ClientDetails{
		Pid:        os.Getpid(),
		ClientType: keybase1.ClientType_KBFS,
		Argv:       os.Args,
		Version:    VersionString(),
	})
	if err != nil {
		return err
	}

	return nil
}

// OnConnectError implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnConnectError(err error, wait time.Duration) {
	k.log.Warning("KeybaseDaemonRPC: connection error: %q; retrying in %s",
		err, wait)
}

// OnDoCommandError implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnDoCommandError(err error, wait time.Duration) {
	k.log.Warning("KeybaseDaemonRPC: docommand error: %q; retrying in %s",
		err, wait)
}

// OnDisconnected implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnDisconnected(_ context.Context,
	status rpc.DisconnectStatus) {
	if status == rpc.StartingNonFirstConnection {
		k.log.Warning("KeybaseDaemonRPC is disconnected")
	}

	k.clearCaches()
}

// ShouldRetry implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) ShouldRetry(rpcName string, err error) bool {
	return false
}

// ShouldRetryOnConnect implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) ShouldRetryOnConnect(err error) bool {
	_, inputCanceled := err.(libkb.InputCanceledError)
	return !inputCanceled
}

func (k *KeybaseDaemonRPC) keepAliveLoop(ctx context.Context) {
	// If the connection is dropped, we need to re-connect and send
	// another HelloIAm message. However, we can't actually detect
	// when the connection is closed until KBFS makes another outgoing
	// RPC, which might not happen for a while.  So continuously send
	// a cheap RPC in the background, so that OnConnect will get
	// called as soon as the connection comes back.
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if k.sessionClient == nil {
				// Clients haven't been filled yet.
				continue
			}
			const sessionID = 0
			err := k.sessionClient.SessionPing(ctx)
			if err != nil {
				k.log.CWarningf(
					ctx, "Background keep alive hit an error: %v", err)
			}
		}
	}
}

// Shutdown implements the KeybaseService interface for KeybaseDaemonRPC.
func (k *KeybaseDaemonRPC) Shutdown() {
	if k.shutdownFn != nil {
		k.shutdownFn()
	}
	if k.keepAliveCancel != nil {
		k.keepAliveCancel()
	}
}
<commit_msg>keybase_daemon_rpc: single-ops shouldn't register as KBFS client<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package libkbfs

import (
	"os"
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	"golang.org/x/net/context"
)

// KeybaseDaemonRPC implements the KeybaseService interface using RPC
// calls.
type KeybaseDaemonRPC struct {
	*KeybaseServiceBase
	daemonLog logger.Logger

	// Only used when there's a real connection (i.e., not in
	// testing).
	shutdownFn func()

	keepAliveCancel context.CancelFunc

	// protocols (additional to required protocols) to register on server connect
	protocols []rpc.Protocol

	// simplefs is the simplefs implementation used (if not nil)
	simplefs keybase1.SimpleFSInterface
}

var _ keybase1.NotifySessionInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyKeyfamilyInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyPaperKeyInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyFSRequestInterface = (*KeybaseDaemonRPC)(nil)

var _ keybase1.NotifyTeamInterface = (*KeybaseDaemonRPC)(nil)

var _ rpc.ConnectionHandler = (*KeybaseDaemonRPC)(nil)

var _ KeybaseService = (*KeybaseDaemonRPC)(nil)

// NewKeybaseDaemonRPC makes a new KeybaseDaemonRPC that makes RPC
// calls using the socket of the given Keybase context.
func NewKeybaseDaemonRPC(config Config, kbCtx Context, log logger.Logger,
	debug bool, createSimpleFS func(Config) keybase1.SimpleFSInterface,
) *KeybaseDaemonRPC {
	k := newKeybaseDaemonRPC(config, kbCtx, log)
	k.config = config
	k.daemonLog = logger.NewWithCallDepth("daemon", 1)
	if createSimpleFS != nil {
		k.simplefs = createSimpleFS(config)
	}
	if debug {
		k.daemonLog.Configure("", true, "")
	}
	conn := NewSharedKeybaseConnection(kbCtx, config, k)
	k.fillClients(conn.GetClient())
	k.shutdownFn = conn.Shutdown

	if config.Mode() != InitMinimal {
		ctx, cancel := context.WithCancel(context.Background())
		k.keepAliveCancel = cancel
		go k.keepAliveLoop(ctx)
	}

	return k
}

// For testing.
func newKeybaseDaemonRPCWithClient(kbCtx Context, client rpc.GenericClient,
	log logger.Logger) *KeybaseDaemonRPC {
	k := newKeybaseDaemonRPC(nil, kbCtx, log)
	k.fillClients(client)
	// No need for a keepalive loop in this case, since this is only
	// used during testing.
	return k
}

func newKeybaseDaemonRPC(config Config, kbCtx Context, log logger.Logger) *KeybaseDaemonRPC {
	serviceBase := NewKeybaseServiceBase(config, kbCtx, log)
	if serviceBase == nil {
		return nil
	}
	k := KeybaseDaemonRPC{
		KeybaseServiceBase: serviceBase,
	}
	return &k
}

func (k *KeybaseDaemonRPC) fillClients(client rpc.GenericClient) {
	k.FillClients(keybase1.IdentifyClient{Cli: client},
		keybase1.UserClient{Cli: client},
		keybase1.TeamsClient{Cli: client},
		keybase1.MerkleClient{Cli: client},
		keybase1.SessionClient{Cli: client},
		keybase1.FavoriteClient{Cli: client},
		keybase1.KbfsClient{Cli: client},
		keybase1.KbfsMountClient{Cli: client})
}

type daemonLogUI struct {
	log logger.Logger
}

var _ keybase1.LogUiInterface = daemonLogUI{}

func (l daemonLogUI) Log(ctx context.Context, arg keybase1.LogArg) error {
	format := "%s"
	// arg.Text.Markup should always be false, so ignore it.
	s := arg.Text.Data
	switch arg.Level {
	case keybase1.LogLevel_DEBUG:
		l.log.CDebugf(ctx, format, s)
	case keybase1.LogLevel_INFO:
		l.log.CInfof(ctx, format, s)
	case keybase1.LogLevel_WARN:
		l.log.CWarningf(ctx, format, s)
	case keybase1.LogLevel_ERROR:
		l.log.CErrorf(ctx, format, s)
	case keybase1.LogLevel_NOTICE:
		l.log.CNoticef(ctx, format, s)
	case keybase1.LogLevel_CRITICAL:
		l.log.CCriticalf(ctx, format, s)
	default:
		l.log.CWarningf(ctx, format, s)
	}
	return nil
}

type daemonIdentifyUI struct {
	log logger.Logger
}

var _ keybase1.IdentifyUiInterface = daemonIdentifyUI{}

func (d daemonIdentifyUI) DelegateIdentifyUI(ctx context.Context) (int, error) {
	d.log.CDebugf(ctx, "DelegateIdentifyUI() (returning 0, UIDelegationUnavailableError{}")
	return 0, libkb.UIDelegationUnavailableError{}
}

func (d daemonIdentifyUI) Start(ctx context.Context, arg keybase1.StartArg) error {
	d.log.CDebugf(ctx, "Start(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayKey(ctx context.Context, arg keybase1.DisplayKeyArg) error {
	d.log.CDebugf(ctx, "DisplayKey(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) ReportLastTrack(ctx context.Context, arg keybase1.ReportLastTrackArg) error {
	d.log.CDebugf(ctx, "ReportLastTrack(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) LaunchNetworkChecks(ctx context.Context, arg keybase1.LaunchNetworkChecksArg) error {
	d.log.CDebugf(ctx, "LaunchNetworkChecks(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayTrackStatement(ctx context.Context, arg keybase1.DisplayTrackStatementArg) error {
	d.log.CDebugf(ctx, "DisplayTrackStatement(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) FinishWebProofCheck(ctx context.Context, arg keybase1.FinishWebProofCheckArg) error {
	d.log.CDebugf(ctx, "FinishWebProofCheck(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) FinishSocialProofCheck(ctx context.Context, arg keybase1.FinishSocialProofCheckArg) error {
	d.log.CDebugf(ctx, "FinishSocialProofCheck(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayCryptocurrency(ctx context.Context, arg keybase1.DisplayCryptocurrencyArg) error {
	d.log.CDebugf(ctx, "DisplayCryptocurrency(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) ReportTrackToken(ctx context.Context, arg keybase1.ReportTrackTokenArg) error {
	d.log.CDebugf(ctx, "ReportTrackToken(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) DisplayUserCard(ctx context.Context, arg keybase1.DisplayUserCardArg) error {
	d.log.CDebugf(ctx, "DisplayUserCard(%+v)", arg)
	return nil
}

func (d daemonIdentifyUI) Confirm(ctx context.Context, arg keybase1.ConfirmArg) (keybase1.ConfirmResult, error) {
	d.log.CDebugf(ctx, "Confirm(%+v) (returning false)", arg)
	return keybase1.ConfirmResult{
		IdentityConfirmed: false,
		RemoteConfirmed:   false,
	}, nil
}

func (d daemonIdentifyUI) DisplayTLFCreateWithInvite(ctx context.Context,
	arg keybase1.DisplayTLFCreateWithInviteArg) (err error) {
	return nil
}

func (d daemonIdentifyUI) Dismiss(ctx context.Context,
	_ keybase1.DismissArg) error {
	return nil
}

func (d daemonIdentifyUI) Cancel(ctx context.Context, sessionID int) error {
	d.log.CDebugf(ctx, "Cancel(%d)", sessionID)
	return nil
}

func (d daemonIdentifyUI) Finish(ctx context.Context, sessionID int) error {
	d.log.CDebugf(ctx, "Finish(%d)", sessionID)
	return nil
}

// HandlerName implements the ConnectionHandler interface.
func (*KeybaseDaemonRPC) HandlerName() string {
	return "KeybaseDaemonRPC"
}

// AddProtocols adds protocols that are registered on server connect
func (k *KeybaseDaemonRPC) AddProtocols(protocols []rpc.Protocol) {
	if protocols == nil {
		return
	}
	if k.protocols != nil {
		k.protocols = append(k.protocols, protocols...)
	} else {
		k.protocols = protocols
	}
}

// OnConnect implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnConnect(ctx context.Context,
	conn *rpc.Connection, rawClient rpc.GenericClient,
	server *rpc.Server) error {

	// Protocols that KBFS requires
	protocols := []rpc.Protocol{
		keybase1.LogUiProtocol(daemonLogUI{k.daemonLog}),
		keybase1.IdentifyUiProtocol(daemonIdentifyUI{k.daemonLog}),
		keybase1.NotifySessionProtocol(k),
		keybase1.NotifyKeyfamilyProtocol(k),
		keybase1.NotifyPaperKeyProtocol(k),
		keybase1.NotifyFSRequestProtocol(k),
		keybase1.NotifyTeamProtocol(k),
		keybase1.TlfKeysProtocol(k),
		keybase1.ReachabilityProtocol(k),
	}

	// Add simplefs if set
	if k.simplefs != nil {
		protocols = append(protocols, keybase1.SimpleFSProtocol(k.simplefs))
	}

	if k.protocols != nil {
		protocols = append(protocols, k.protocols...)
	}

	for _, p := range protocols {
		err := server.Register(p)
		if err != nil {
			if _, ok := err.(rpc.AlreadyRegisteredError); !ok {
				return err
			}
		}
	}

	// Using conn.GetClient() here would cause problematic
	// recursion.
	c := keybase1.NotifyCtlClient{Cli: rawClient}
	err := c.SetNotifications(ctx, keybase1.NotificationChannels{
		Session:      true,
		Paperkeys:    true,
		Keyfamily:    true,
		Kbfsrequest:  true,
		Reachability: true,
		Team:         true,
	})
	if err != nil {
		return err
	}

	// Introduce ourselves. TODO: move this to SharedKeybaseConnection
	// somehow?
	configClient := keybase1.ConfigClient{Cli: rawClient}
	ct := keybase1.ClientType_KBFS
	if k.config.Mode() == InitSingleOp {
		ct = keybase1.ClientType_NONE
	}
	err = configClient.HelloIAm(ctx, keybase1.ClientDetails{
		Pid:        os.Getpid(),
		ClientType: ct,
		Argv:       os.Args,
		Version:    VersionString(),
	})
	if err != nil {
		return err
	}

	return nil
}

// OnConnectError implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnConnectError(err error, wait time.Duration) {
	k.log.Warning("KeybaseDaemonRPC: connection error: %q; retrying in %s",
		err, wait)
}

// OnDoCommandError implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnDoCommandError(err error, wait time.Duration) {
	k.log.Warning("KeybaseDaemonRPC: docommand error: %q; retrying in %s",
		err, wait)
}

// OnDisconnected implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) OnDisconnected(_ context.Context,
	status rpc.DisconnectStatus) {
	if status == rpc.StartingNonFirstConnection {
		k.log.Warning("KeybaseDaemonRPC is disconnected")
	}

	k.clearCaches()
}

// ShouldRetry implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) ShouldRetry(rpcName string, err error) bool {
	return false
}

// ShouldRetryOnConnect implements the ConnectionHandler interface.
func (k *KeybaseDaemonRPC) ShouldRetryOnConnect(err error) bool {
	_, inputCanceled := err.(libkb.InputCanceledError)
	return !inputCanceled
}

func (k *KeybaseDaemonRPC) keepAliveLoop(ctx context.Context) {
	// If the connection is dropped, we need to re-connect and send
	// another HelloIAm message. However, we can't actually detect
	// when the connection is closed until KBFS makes another outgoing
	// RPC, which might not happen for a while.  So continuously send
	// a cheap RPC in the background, so that OnConnect will get
	// called as soon as the connection comes back.
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if k.sessionClient == nil {
				// Clients haven't been filled yet.
				continue
			}
			const sessionID = 0
			err := k.sessionClient.SessionPing(ctx)
			if err != nil {
				k.log.CWarningf(
					ctx, "Background keep alive hit an error: %v", err)
			}
		}
	}
}

// Shutdown implements the KeybaseService interface for KeybaseDaemonRPC.
func (k *KeybaseDaemonRPC) Shutdown() {
	if k.shutdownFn != nil {
		k.shutdownFn()
	}
	if k.keepAliveCancel != nil {
		k.keepAliveCancel()
	}
}
<|endoftext|>


<commit_before>package nats

import (
	"net"
	"sync"
)

type Subscription struct {
	sr *subscriptionRegistry

	sid      uint
	frozen   bool
	maximum  uint
	received uint

	subject string
	queue   string

	Inbox chan *readMessage
}

func (s *Subscription) freeze() {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.frozen = true
}

func (s *Subscription) SetSubject(v string) {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.subject = v
}

func (s *Subscription) SetQueue(v string) {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.queue = v
}

func (s *Subscription) SetMaximum(v uint) {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.maximum = v
}

func (s *Subscription) writeSubscribe() writeObject {
	var o = new(writeSubscribe)

	o.Sid = s.sid
	o.Subject = s.subject
	o.Queue = s.queue

	return o
}

func (s *Subscription) writeUnsubscribe(includeMaximum bool) writeObject {
	var o = new(writeUnsubscribe)

	o.Sid = s.sid

	if includeMaximum {
		o.Maximum = s.maximum
	}

	return o
}

func (s *Subscription) Subscribe() {
	s.sr.Subscribe(s)
}

func (s *Subscription) Unsubscribe() {
	s.sr.Unsubscribe(s)
}

func (s *Subscription) deliver(m *readMessage) {
	s.received++
	s.Inbox <- m

	// Unsubscribe if the maximum number of messages has been received
	if s.maximum > 0 && s.received >= s.maximum {
		s.Unsubscribe()
	}
}

type subscriptionRegistry struct {
	sync.Mutex
	*Client

	sid uint
	m   map[uint]*Subscription
}

func (sr *subscriptionRegistry) emptyMap() {
	sr.m = make(map[uint]*Subscription)
}

func (sr *subscriptionRegistry) setup(c *Client) {
	sr.Client = c

	sr.emptyMap()
}

func (sr *subscriptionRegistry) teardown() {
	sr.Lock()
	defer sr.Unlock()

	for _, s := range sr.m {
		close(s.Inbox)
	}

	sr.emptyMap()
}

func (sr *subscriptionRegistry) NewSubscription(sub string) *Subscription {
	var s = new(Subscription)

	sr.Lock()

	s.sr = sr

	sr.sid++
	s.sid = sr.sid

	sr.Unlock()

	s.SetSubject(sub)
	s.Inbox = make(chan *readMessage)

	return s
}

func (sr *subscriptionRegistry) Subscribe(s *Subscription) {
	sr.Lock()

	sr.m[s.sid] = s
	s.freeze()

	sr.Unlock()

	sr.Client.Write(s.writeSubscribe())

	if s.maximum > 0 {
		sr.Client.Write(s.writeUnsubscribe(true))
	}

	return
}

func (sr *subscriptionRegistry) Unsubscribe(s *Subscription) {
	sr.Lock()

	delete(sr.m, s.sid)

	// Since this subscription is now removed from the registry, it will no
	// longer receive messages and the inbox can be closed
	close(s.Inbox)

	sr.Unlock()

	sr.Client.Write(s.writeUnsubscribe(false))

	return
}

func (sr *subscriptionRegistry) deliver(m *readMessage) {
	var s *Subscription
	var ok bool

	sr.Lock()
	s, ok = sr.m[m.SubscriptionId]
	sr.Unlock()

	if ok {
		s.deliver(m)
	}
}

type Client struct {
	subscriptionRegistry
	Handshaker

	Addr string
	User string
	Pass string

	cc chan *Connection

	// Notify running client to stop
	sc chan bool
}

func NewClient(addr string) *Client {
	var t = new(Client)

	t.subscriptionRegistry.setup(t)

	t.Handshaker = &Handshake{}

	t.Addr = addr

	t.cc = make(chan *Connection)

	t.sc = make(chan bool)

	return t
}

func (t *Client) Write(o writeObject) bool {
	var c *Connection
	var ok bool

	c, ok = <-t.cc
	if !ok {
		return false
	}

	return c.Write(o)
}

func (t *Client) Ping() bool {
	var c *Connection
	var ok bool

	c, ok = <-t.cc
	if !ok {
		return false
	}

	return c.Ping()
}

func (t *Client) publish(s string, m []byte, confirm bool) bool {
	var o = new(writePublish)
	var c *Connection
	var ok bool

	o.Subject = s
	o.Message = m

	c, ok = <-t.cc
	if !ok {
		return false
	}

	ok = c.Write(o)
	if !ok {
		return false
	}

	// Round trip to confirm the publish was received
	if confirm {
		return c.Ping()
	}

	return true
}

func (t *Client) Publish(s string, m []byte) bool {
	return t.publish(s, m, false)
}

func (t *Client) PublishAndConfirm(s string, m []byte) bool {
	return t.publish(s, m, true)
}

func (t *Client) Stop() {
	t.sc <- true
}

func (t *Client) runConnection(n net.Conn) error {
	var e error
	var c *Connection
	var dc chan bool

	c = NewConnection(n)
	dc = make(chan bool)

	// Feed connection until stop
	go func() {
		for {
			select {
			case t.cc <- c:
			case <-t.sc:
				c.Stop()

				// Wait for c.Run() to return and notify dc
				<-dc
				return

			case <-dc:
				return
			}
		}
	}()

	// Read messages until EOF
	go func() {
		var o readObject

		for o = range c.oc {
			switch oo := o.(type) {
			case *readMessage:
				t.deliver(oo)
			}
		}
	}()

	e = c.Run()
	dc <- true
	close(t.cc)

	return e
}

func (t *Client) Run(n net.Conn) error {
	var e error

	n, e = t.Handshake(n)
	if e != nil {
		return e
	}

	e = t.runConnection(n)
	if e != nil {
		return e
	}

	// No error means the client was stopped, and we should close all subscriptions
	t.subscriptionRegistry.teardown()

	return nil
}
<commit_msg>User and pass don't belong in client struct<commit_after>package nats

import (
	"net"
	"sync"
)

type Subscription struct {
	sr *subscriptionRegistry

	sid      uint
	frozen   bool
	maximum  uint
	received uint

	subject string
	queue   string

	Inbox chan *readMessage
}

func (s *Subscription) freeze() {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.frozen = true
}

func (s *Subscription) SetSubject(v string) {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.subject = v
}

func (s *Subscription) SetQueue(v string) {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.queue = v
}

func (s *Subscription) SetMaximum(v uint) {
	if s.frozen {
		panic("subscription is frozen")
	}

	s.maximum = v
}

func (s *Subscription) writeSubscribe() writeObject {
	var o = new(writeSubscribe)

	o.Sid = s.sid
	o.Subject = s.subject
	o.Queue = s.queue

	return o
}

func (s *Subscription) writeUnsubscribe(includeMaximum bool) writeObject {
	var o = new(writeUnsubscribe)

	o.Sid = s.sid

	if includeMaximum {
		o.Maximum = s.maximum
	}

	return o
}

func (s *Subscription) Subscribe() {
	s.sr.Subscribe(s)
}

func (s *Subscription) Unsubscribe() {
	s.sr.Unsubscribe(s)
}

func (s *Subscription) deliver(m *readMessage) {
	s.received++
	s.Inbox <- m

	// Unsubscribe if the maximum number of messages has been received
	if s.maximum > 0 && s.received >= s.maximum {
		s.Unsubscribe()
	}
}

type subscriptionRegistry struct {
	sync.Mutex
	*Client

	sid uint
	m   map[uint]*Subscription
}

func (sr *subscriptionRegistry) emptyMap() {
	sr.m = make(map[uint]*Subscription)
}

func (sr *subscriptionRegistry) setup(c *Client) {
	sr.Client = c

	sr.emptyMap()
}

func (sr *subscriptionRegistry) teardown() {
	sr.Lock()
	defer sr.Unlock()

	for _, s := range sr.m {
		close(s.Inbox)
	}

	sr.emptyMap()
}

func (sr *subscriptionRegistry) NewSubscription(sub string) *Subscription {
	var s = new(Subscription)

	sr.Lock()

	s.sr = sr

	sr.sid++
	s.sid = sr.sid

	sr.Unlock()

	s.SetSubject(sub)
	s.Inbox = make(chan *readMessage)

	return s
}

func (sr *subscriptionRegistry) Subscribe(s *Subscription) {
	sr.Lock()

	sr.m[s.sid] = s
	s.freeze()

	sr.Unlock()

	sr.Client.Write(s.writeSubscribe())

	if s.maximum > 0 {
		sr.Client.Write(s.writeUnsubscribe(true))
	}

	return
}

func (sr *subscriptionRegistry) Unsubscribe(s *Subscription) {
	sr.Lock()

	delete(sr.m, s.sid)

	// Since this subscription is now removed from the registry, it will no
	// longer receive messages and the inbox can be closed
	close(s.Inbox)

	sr.Unlock()

	sr.Client.Write(s.writeUnsubscribe(false))

	return
}

func (sr *subscriptionRegistry) deliver(m *readMessage) {
	var s *Subscription
	var ok bool

	sr.Lock()
	s, ok = sr.m[m.SubscriptionId]
	sr.Unlock()

	if ok {
		s.deliver(m)
	}
}

type Client struct {
	subscriptionRegistry
	Handshaker

	Addr string

	cc chan *Connection

	// Notify running client to stop
	sc chan bool
}

func NewClient(addr string) *Client {
	var t = new(Client)

	t.subscriptionRegistry.setup(t)

	t.Handshaker = &Handshake{}

	t.Addr = addr

	t.cc = make(chan *Connection)

	t.sc = make(chan bool)

	return t
}

func (t *Client) Write(o writeObject) bool {
	var c *Connection
	var ok bool

	c, ok = <-t.cc
	if !ok {
		return false
	}

	return c.Write(o)
}

func (t *Client) Ping() bool {
	var c *Connection
	var ok bool

	c, ok = <-t.cc
	if !ok {
		return false
	}

	return c.Ping()
}

func (t *Client) publish(s string, m []byte, confirm bool) bool {
	var o = new(writePublish)
	var c *Connection
	var ok bool

	o.Subject = s
	o.Message = m

	c, ok = <-t.cc
	if !ok {
		return false
	}

	ok = c.Write(o)
	if !ok {
		return false
	}

	// Round trip to confirm the publish was received
	if confirm {
		return c.Ping()
	}

	return true
}

func (t *Client) Publish(s string, m []byte) bool {
	return t.publish(s, m, false)
}

func (t *Client) PublishAndConfirm(s string, m []byte) bool {
	return t.publish(s, m, true)
}

func (t *Client) Stop() {
	t.sc <- true
}

func (t *Client) runConnection(n net.Conn) error {
	var e error
	var c *Connection
	var dc chan bool

	c = NewConnection(n)
	dc = make(chan bool)

	// Feed connection until stop
	go func() {
		for {
			select {
			case t.cc <- c:
			case <-t.sc:
				c.Stop()

				// Wait for c.Run() to return and notify dc
				<-dc
				return

			case <-dc:
				return
			}
		}
	}()

	// Read messages until EOF
	go func() {
		var o readObject

		for o = range c.oc {
			switch oo := o.(type) {
			case *readMessage:
				t.deliver(oo)
			}
		}
	}()

	e = c.Run()
	dc <- true
	close(t.cc)

	return e
}

func (t *Client) Run(n net.Conn) error {
	var e error

	n, e = t.Handshake(n)
	if e != nil {
		return e
	}

	e = t.runConnection(n)
	if e != nil {
		return e
	}

	// No error means the client was stopped, and we should close all subscriptions
	t.subscriptionRegistry.teardown()

	return nil
}
<|endoftext|>