<commit_before>// Copyright (c) 2016 Company 0, LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

// rpc contains all structures required by the ZK protocol.
//
// A ZK session has two discrete phases:
//	1. pre session phase, used to create accounts and obtain zkserver key
//	2. session phase, used for all other RPC commands
//	3. once the key exchange is complete the server shall issue a Welcome
//         command.  The welcome command also transfer additional settings such
//         as tag depth etc.
//
// In order to exchange messages with a third party two pieces of information
// are required.  Each side must know the other's long lived public identity
// and the public DH ratchet keys.
// The process, using RPC, to obtains that information is as follows:
//	1. Alice sends Bob a Rendezvous command that contains her encrypted
//	   identity. She uses a third party communication method (phone, IRC
//	   etc)	to share the rendezvous PIN code and a shared password.
//	2. Bob obtains Alice's identity by sending a RendezvousPull command
//	   using the PIN code.  After decrypting Alice's identity blob using the
//	   share password he replies with a Cache command that contains his long
//	   lived public identity and his initial public DH ratchet keys.
//	3. Alice is notified, using the normal Push RPC mechanism, when Bob has
//	   replied.  She then replies to Bob with her public DH ratchet keys.
//
// The external identity and key exchange process is outside of the scope of
// this document.
package rpc

import (
	"crypto/sha256"
	"errors"
	"strconv"

	"github.com/companyzero/zkc/ratchet"
	"github.com/companyzero/zkc/zkidentity"
)

const (
	// pre session phase
	InitialCmdIdentify      = "identify"
	InitialCmdCreateAccount = "createaccount"
	InitialCmdSession       = "session"

	// session phase
	SessionCmdWelcome = "welcome"

	// tagged server commands
	TaggedCmdRendezvous          = "rendezvous"
	TaggedCmdRendezvousReply     = "rendezvousreply"
	TaggedCmdRendezvousPull      = "rendezvouspull"
	TaggedCmdRendezvousPullReply = "rendezvouspullreply"
	TaggedCmdCache               = "cache"
	TaggedCmdPush                = "push"
	TaggedCmdAcknowledge         = "ack"
	TaggedCmdPing                = "ping"
	TaggedCmdPong                = "pong"

	// misc
	MessageModeNormal = 0
	MessageModeMe     = 1
)

// CreateAccount is a PRPC that is used to create a new account on the server.
// Policy dictates if this is allowed or not.
type CreateAccount struct {
	Token          string                    // auth token
	PublicIdentity zkidentity.PublicIdentity // long lived public identity
}

// sanitized errors for CreateAccountReply
var (
	ErrCreateDisallowed = errors.New("not allowed")
	ErrInternalError    = errors.New("internal error, contact administrator")
)

// Message is the generic command that flows between a server and client and
// vice versa.  Its purpose is to add a discriminator to simplify payload
// decoding.  Additionally it has a tag that the recipient shall return
// unmodified when replying.  The tag is originated by the sender and shall be
// unique provided an answer is expected.  The receiver shall not interpret or
// use the tag in any way.
type Message struct {
	Command   string // discriminator
	TimeStamp int64  // originator timestamp
	Tag       uint32 // client generated tag, shall be unique
	//followed by Payload []byte
}

// Empty is used when there is no Payload.
type Empty struct{}

const (
	ProtocolVersion = 5
)

// Welcome is written immediately following a key exchange.  This command
// purpose is to detect if the key exchange completed on the client side.  If
// the key exchange failed the server will simply disconnect.
type Welcome struct {
	Version    int   // protocol version
	ServerTime int64 // server timestamp

	// Client shall ensure it is compatible with the server requirements
	Properties []ServerProperty // server properties
}

type ServerProperty struct {
	Key      string // name of property
	Value    string // value of property
	Required bool   // if true client must handle this entry
}

const (
	// Tag Depth is a required property.  It defines maximum outstanding
	// commands.
	PropTagDepth        = "tagdepth"
	PropTagDepthDefault = "10"

	// MOTD (Message Of The Day) is an optional property.  It is a welcome
	// message that is sent from the server to the client upon first
	// contact.  The client may display this.
	PropMOTD = "motd"

	// Max Attachment Size is a required property.  It defines the maximum
	// attachment size.  Attachment size is defined as the largest size a
	// file transfer is allowed to be.
	PropMaxAttachmentSize        = "maxattachmentsize"
	PropMaxAttachmentSizeDefault = uint64(10 * 1024 * 1024)

	// Max Chunk Size is a required property.  It defines the maximum chunk
	// size.  Chunk size is defined as the largest size a CRPC is allowed
	// to be.
	PropMaxChunkSize        = "maxchunksize"
	PropMaxChunkSizeDefault = uint64(256 * 1024)

	// Max Message Size is a required property.  It defines the maximum
	// message size.  Message size is defined as the largest size a CRPC is
	// allowed to be.  This includes message overhead etc.
	PropMaxMsgSize        = "maxmsgsize"
	PropMaxMsgSizeDefault = PropMaxChunkSizeDefault + 1024

	// Server Time is a required property.  It contains the server time
	// stamp.  The client shall warn the user if the client is not time
	// synced.  Clients and proxies really shall run NTP.
	PropServerTime = "servertime"
)

var (
	// required
	DefaultPropTagDepth = ServerProperty{
		Key:      PropTagDepth,
		Value:    PropTagDepthDefault,
		Required: true,
	}
	DefaultPropMaxAttachmentSize = ServerProperty{
		Key:      PropMaxAttachmentSize,
		Value:    strconv.FormatUint(PropMaxAttachmentSizeDefault, 10),
		Required: true,
	}
	DefaultPropMaxChunkSize = ServerProperty{
		Key:      PropMaxChunkSize,
		Value:    strconv.FormatUint(PropMaxChunkSizeDefault, 10),
		Required: true,
	}
	DefaultPropMaxMsgSize = ServerProperty{
		Key:      PropMaxMsgSize,
		Value:    strconv.FormatUint(PropMaxMsgSizeDefault, 10),
		Required: true,
	}
	DefaultServerTime = ServerProperty{
		Key:      PropServerTime,
		Value:    "", // int64 unix time
		Required: true,
	}

	// optional
	DefaultPropMOTD = ServerProperty{
		Key:      PropMOTD,
		Value:    "",
		Required: false,
	}

	// All properties must exist in this array.
	SupportedServerProperties = []ServerProperty{
		// required
		DefaultPropTagDepth,
		DefaultPropMaxAttachmentSize,
		DefaultPropMaxChunkSize,
		DefaultPropMaxMsgSize,
		DefaultServerTime,

		// optional
		DefaultPropMOTD,
	}
)

// CreateAccountReply returns a sanitized error to the client indicating
// success or failure of the CreateAccountReply command.  Errors is set to ""
// on success.
type CreateAccountReply struct {
	Error string // if create account failed error contains the reason.
}

// Push is a PRPC that is used to push cached encrypted blobs to a user.  This
// command must be acknowledged by the remote side.
type Push struct {
	From     [32]byte // sender identity
	Received int64    // server received timestamp
	Payload  []byte   // encrypted payload
}

// Cache is a PRPC that is used to store message on server for later push
// delivery.  This command must be acknowledged by the remote side.
type Cache struct {
	To      [32]byte // recipient identity
	Payload []byte   // encrypted payload
}

// Ping is a PRPC that is used to determine if the server is alive.
// This command must be acknowledged by the remote side.
type Ping struct{}
type Pong struct{}

// client to client commands

// Passthrough is a command that flows from a client through the server to
// another client.  Payload contains an encrypted CRPC. These commands shall
// not be cached by the server.
type Passthrough struct {
	To      [32]byte // recipient identity
	Payload []byte   // encrypted payload
}

// Rendezvous sends a blob to the server. Blob shall be < 4096 and and
// expiration shall be < 168 (7 * 24).
type Rendezvous struct {
	Blob       []byte // data being shared
	Expiration string // hours until Rendezvous expires
}

// RendezvousReply is a reply packet for a Rendezvous command.  Token contains
// an easy to remember PIN code to identify initial Rendezvous blob.
type RendezvousReply struct {
	Token string // Rendezvous token that identifies blob
	Error string // If an error occured Error will be != ""
}

// RendezvousPull tries to download a previously uploaded blob.
type RendezvousPull struct {
	Token string // Rendezvous token that identifies blob
}

// RendezvousPullReply contains a data blob reply to a previous RendezvousPull
// command that is identified by token.
type RendezvousPullReply struct {
	Error string // set if an error occurred
	Token string // Rendezvous token that identifies blob
	Blob  []byte // data reply to previous Rendezvous
}

// IdentityKX contains the long lived public identify and the DH ratchet keys.
// It is the second step during the IDKX exchange.
type IdentityKX struct {
	Identity zkidentity.PublicIdentity
	KX       ratchet.KeyExchange
}

// KX contains the DH ratchet keys.  It is the third step during the IDKX
// exchange.
type KX struct {
	KX ratchet.KeyExchange
}

const (
	// CRPC commands
	CRPCCmdPrivateMessage = "privmsg"
	CRPCCmdGroupInvite    = "groupinvite"
	CRPCCmdGroupJoin      = "groupjoin"
	CRPCCmdGroupPart      = "grouppart"
	CRPCCmdGroupKill      = "groupkill"
	CRPCCmdGroupKick      = "groupkick"
	CRPCCmdGroupList      = "grouplist"
	CRPCCmdGroupListGet   = "grouplistget" // no struct used
	CRPCCmdGroupMessage   = "groupmessage"
	CRPCCmdChunkNew       = "chunknew"
	CRPCCmdChunk          = "chunk"
	CRPCCmdJanitorMessage = "janitormessage"

	// compression
	CRPCCompNone = ""
	CRPCCompZLIB = "zlib"

	// janitor
	CRPCJanitorRatchetReset = "ratchetreset"
	CRPCJanitorDeleted      = "deleted"
)

// CRPC is a client RPC message.
type CRPC struct {
	Timestamp   int64  // client side timestamp
	Command     string // discriminator
	Compression string // compression used on Payload
	//followed by Payload []byte
}

// PrivateMessage is a CRPC that contains a text message.
type PrivateMessage struct {
	Text string
	Mode uint32 // 0 regular mode, 1 /me
}

// JanitorMessage is a CRPC that tells the other party some sort of
// housekeeping occurred.
type JanitorMessage struct {
	Command string
	Reason  string
}

// GroupInvite, sender is implicit to CRPC.
// XXX Note that there is no explicit way to prohibit sender being admin.
// XXX This needs some more thought.
type GroupInvite struct {
	Name        string // group name
	Token       uint64 // invite token
	Description string // group description
	Expires     int64  // unix time when this invite expires
}

// GroupJoin
type GroupJoin struct {
	Name  string // group name
	Token uint64 // invite token, implicitly identifies sender
	Error string // accept or deny Invite
}

// GroupPart, sender is implicit to CRPC
type GroupPart struct {
	Name   string // group name
	Reason string // reason to depart group
}

// GroupKill, sender is implicit to CRPC
type GroupKill struct {
	Name   string // group name
	Reason string // reason to disassemble group
}

// GroupKick, sender is implicit to CRPC
type GroupKick struct {
	Member       [zkidentity.IdentitySize]byte // kickee
	Reason       string                        // why member was kicked
	Parted       bool                          // kicked/parted
	NewGroupList GroupList                     // new GroupList
}

// GroupList, currently we detect spoofing by ensuring the origin of the
// message.  This may not be sufficient and we may have to add a signature of
// sorts.  For now roll with this assumption.
type GroupList struct {
	Name       string // group name
	Generation uint64 // incremented every time list changes
	Timestamp  int64  // unix time last generation changed

	// all participants, [0] is administrator
	// receiver must check [0] == originator
	Members [][zkidentity.IdentitySize]byte
}

// GroupMessage is a message to a group.
type GroupMessage struct {
	Name       string // group name
	Generation uint64 // Generation used
	Message    string // Actual message
	Mode       uint32 // 0 regular mode, 1 /me
}

// ChunkNew describes a chunked file transfer initiation.
type ChunkNew struct {
	Size        uint64            // total file size
	ChunkSize   uint64            // chunk size
	Filename    string            // original filename
	Description string            // user provided description
	MIME        string            // mime type
	Digest      [sha256.Size]byte // digest of file -> unique identifier
}

type Chunk struct {
	Offset  uint64            // offset in file
	Digest  [sha256.Size]byte // digest of file -> unique identifier
	Payload []byte            // chunk
}
<commit_msg>crank the protocol version; reminded by marco@<commit_after>// Copyright (c) 2016 Company 0, LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

// rpc contains all structures required by the ZK protocol.
//
// A ZK session has two discrete phases:
//	1. pre session phase, used to create accounts and obtain zkserver key
//	2. session phase, used for all other RPC commands
//	3. once the key exchange is complete the server shall issue a Welcome
//         command.  The welcome command also transfer additional settings such
//         as tag depth etc.
//
// In order to exchange messages with a third party two pieces of information
// are required.  Each side must know the other's long lived public identity
// and the public DH ratchet keys.
// The process, using RPC, to obtains that information is as follows:
//	1. Alice sends Bob a Rendezvous command that contains her encrypted
//	   identity. She uses a third party communication method (phone, IRC
//	   etc)	to share the rendezvous PIN code and a shared password.
//	2. Bob obtains Alice's identity by sending a RendezvousPull command
//	   using the PIN code.  After decrypting Alice's identity blob using the
//	   share password he replies with a Cache command that contains his long
//	   lived public identity and his initial public DH ratchet keys.
//	3. Alice is notified, using the normal Push RPC mechanism, when Bob has
//	   replied.  She then replies to Bob with her public DH ratchet keys.
//
// The external identity and key exchange process is outside of the scope of
// this document.
package rpc

import (
	"crypto/sha256"
	"errors"
	"strconv"

	"github.com/companyzero/zkc/ratchet"
	"github.com/companyzero/zkc/zkidentity"
)

const (
	// pre session phase
	InitialCmdIdentify      = "identify"
	InitialCmdCreateAccount = "createaccount"
	InitialCmdSession       = "session"

	// session phase
	SessionCmdWelcome = "welcome"

	// tagged server commands
	TaggedCmdRendezvous          = "rendezvous"
	TaggedCmdRendezvousReply     = "rendezvousreply"
	TaggedCmdRendezvousPull      = "rendezvouspull"
	TaggedCmdRendezvousPullReply = "rendezvouspullreply"
	TaggedCmdCache               = "cache"
	TaggedCmdPush                = "push"
	TaggedCmdAcknowledge         = "ack"
	TaggedCmdPing                = "ping"
	TaggedCmdPong                = "pong"

	// misc
	MessageModeNormal = 0
	MessageModeMe     = 1
)

// CreateAccount is a PRPC that is used to create a new account on the server.
// Policy dictates if this is allowed or not.
type CreateAccount struct {
	Token          string                    // auth token
	PublicIdentity zkidentity.PublicIdentity // long lived public identity
}

// sanitized errors for CreateAccountReply
var (
	ErrCreateDisallowed = errors.New("not allowed")
	ErrInternalError    = errors.New("internal error, contact administrator")
)

// Message is the generic command that flows between a server and client and
// vice versa.  Its purpose is to add a discriminator to simplify payload
// decoding.  Additionally it has a tag that the recipient shall return
// unmodified when replying.  The tag is originated by the sender and shall be
// unique provided an answer is expected.  The receiver shall not interpret or
// use the tag in any way.
type Message struct {
	Command   string // discriminator
	TimeStamp int64  // originator timestamp
	Tag       uint32 // client generated tag, shall be unique
	//followed by Payload []byte
}

// Empty is used when there is no Payload.
type Empty struct{}

const (
	ProtocolVersion = 6
)

// Welcome is written immediately following a key exchange.  This command
// purpose is to detect if the key exchange completed on the client side.  If
// the key exchange failed the server will simply disconnect.
type Welcome struct {
	Version    int   // protocol version
	ServerTime int64 // server timestamp

	// Client shall ensure it is compatible with the server requirements
	Properties []ServerProperty // server properties
}

type ServerProperty struct {
	Key      string // name of property
	Value    string // value of property
	Required bool   // if true client must handle this entry
}

const (
	// Tag Depth is a required property.  It defines maximum outstanding
	// commands.
	PropTagDepth        = "tagdepth"
	PropTagDepthDefault = "10"

	// MOTD (Message Of The Day) is an optional property.  It is a welcome
	// message that is sent from the server to the client upon first
	// contact.  The client may display this.
	PropMOTD = "motd"

	// Max Attachment Size is a required property.  It defines the maximum
	// attachment size.  Attachment size is defined as the largest size a
	// file transfer is allowed to be.
	PropMaxAttachmentSize        = "maxattachmentsize"
	PropMaxAttachmentSizeDefault = uint64(10 * 1024 * 1024)

	// Max Chunk Size is a required property.  It defines the maximum chunk
	// size.  Chunk size is defined as the largest size a CRPC is allowed
	// to be.
	PropMaxChunkSize        = "maxchunksize"
	PropMaxChunkSizeDefault = uint64(256 * 1024)

	// Max Message Size is a required property.  It defines the maximum
	// message size.  Message size is defined as the largest size a CRPC is
	// allowed to be.  This includes message overhead etc.
	PropMaxMsgSize        = "maxmsgsize"
	PropMaxMsgSizeDefault = PropMaxChunkSizeDefault + 1024

	// Server Time is a required property.  It contains the server time
	// stamp.  The client shall warn the user if the client is not time
	// synced.  Clients and proxies really shall run NTP.
	PropServerTime = "servertime"
)

var (
	// required
	DefaultPropTagDepth = ServerProperty{
		Key:      PropTagDepth,
		Value:    PropTagDepthDefault,
		Required: true,
	}
	DefaultPropMaxAttachmentSize = ServerProperty{
		Key:      PropMaxAttachmentSize,
		Value:    strconv.FormatUint(PropMaxAttachmentSizeDefault, 10),
		Required: true,
	}
	DefaultPropMaxChunkSize = ServerProperty{
		Key:      PropMaxChunkSize,
		Value:    strconv.FormatUint(PropMaxChunkSizeDefault, 10),
		Required: true,
	}
	DefaultPropMaxMsgSize = ServerProperty{
		Key:      PropMaxMsgSize,
		Value:    strconv.FormatUint(PropMaxMsgSizeDefault, 10),
		Required: true,
	}
	DefaultServerTime = ServerProperty{
		Key:      PropServerTime,
		Value:    "", // int64 unix time
		Required: true,
	}

	// optional
	DefaultPropMOTD = ServerProperty{
		Key:      PropMOTD,
		Value:    "",
		Required: false,
	}

	// All properties must exist in this array.
	SupportedServerProperties = []ServerProperty{
		// required
		DefaultPropTagDepth,
		DefaultPropMaxAttachmentSize,
		DefaultPropMaxChunkSize,
		DefaultPropMaxMsgSize,
		DefaultServerTime,

		// optional
		DefaultPropMOTD,
	}
)

// CreateAccountReply returns a sanitized error to the client indicating
// success or failure of the CreateAccountReply command.  Errors is set to ""
// on success.
type CreateAccountReply struct {
	Error string // if create account failed error contains the reason.
}

// Push is a PRPC that is used to push cached encrypted blobs to a user.  This
// command must be acknowledged by the remote side.
type Push struct {
	From     [32]byte // sender identity
	Received int64    // server received timestamp
	Payload  []byte   // encrypted payload
}

// Cache is a PRPC that is used to store message on server for later push
// delivery.  This command must be acknowledged by the remote side.
type Cache struct {
	To      [32]byte // recipient identity
	Payload []byte   // encrypted payload
}

// Ping is a PRPC that is used to determine if the server is alive.
// This command must be acknowledged by the remote side.
type Ping struct{}
type Pong struct{}

// client to client commands

// Passthrough is a command that flows from a client through the server to
// another client.  Payload contains an encrypted CRPC. These commands shall
// not be cached by the server.
type Passthrough struct {
	To      [32]byte // recipient identity
	Payload []byte   // encrypted payload
}

// Rendezvous sends a blob to the server. Blob shall be < 4096 and and
// expiration shall be < 168 (7 * 24).
type Rendezvous struct {
	Blob       []byte // data being shared
	Expiration string // hours until Rendezvous expires
}

// RendezvousReply is a reply packet for a Rendezvous command.  Token contains
// an easy to remember PIN code to identify initial Rendezvous blob.
type RendezvousReply struct {
	Token string // Rendezvous token that identifies blob
	Error string // If an error occured Error will be != ""
}

// RendezvousPull tries to download a previously uploaded blob.
type RendezvousPull struct {
	Token string // Rendezvous token that identifies blob
}

// RendezvousPullReply contains a data blob reply to a previous RendezvousPull
// command that is identified by token.
type RendezvousPullReply struct {
	Error string // set if an error occurred
	Token string // Rendezvous token that identifies blob
	Blob  []byte // data reply to previous Rendezvous
}

// IdentityKX contains the long lived public identify and the DH ratchet keys.
// It is the second step during the IDKX exchange.
type IdentityKX struct {
	Identity zkidentity.PublicIdentity
	KX       ratchet.KeyExchange
}

// KX contains the DH ratchet keys.  It is the third step during the IDKX
// exchange.
type KX struct {
	KX ratchet.KeyExchange
}

const (
	// CRPC commands
	CRPCCmdPrivateMessage = "privmsg"
	CRPCCmdGroupInvite    = "groupinvite"
	CRPCCmdGroupJoin      = "groupjoin"
	CRPCCmdGroupPart      = "grouppart"
	CRPCCmdGroupKill      = "groupkill"
	CRPCCmdGroupKick      = "groupkick"
	CRPCCmdGroupList      = "grouplist"
	CRPCCmdGroupListGet   = "grouplistget" // no struct used
	CRPCCmdGroupMessage   = "groupmessage"
	CRPCCmdChunkNew       = "chunknew"
	CRPCCmdChunk          = "chunk"
	CRPCCmdJanitorMessage = "janitormessage"

	// compression
	CRPCCompNone = ""
	CRPCCompZLIB = "zlib"

	// janitor
	CRPCJanitorRatchetReset = "ratchetreset"
	CRPCJanitorDeleted      = "deleted"
)

// CRPC is a client RPC message.
type CRPC struct {
	Timestamp   int64  // client side timestamp
	Command     string // discriminator
	Compression string // compression used on Payload
	//followed by Payload []byte
}

// PrivateMessage is a CRPC that contains a text message.
type PrivateMessage struct {
	Text string
	Mode uint32 // 0 regular mode, 1 /me
}

// JanitorMessage is a CRPC that tells the other party some sort of
// housekeeping occurred.
type JanitorMessage struct {
	Command string
	Reason  string
}

// GroupInvite, sender is implicit to CRPC.
// XXX Note that there is no explicit way to prohibit sender being admin.
// XXX This needs some more thought.
type GroupInvite struct {
	Name        string // group name
	Token       uint64 // invite token
	Description string // group description
	Expires     int64  // unix time when this invite expires
}

// GroupJoin
type GroupJoin struct {
	Name  string // group name
	Token uint64 // invite token, implicitly identifies sender
	Error string // accept or deny Invite
}

// GroupPart, sender is implicit to CRPC
type GroupPart struct {
	Name   string // group name
	Reason string // reason to depart group
}

// GroupKill, sender is implicit to CRPC
type GroupKill struct {
	Name   string // group name
	Reason string // reason to disassemble group
}

// GroupKick, sender is implicit to CRPC
type GroupKick struct {
	Member       [zkidentity.IdentitySize]byte // kickee
	Reason       string                        // why member was kicked
	Parted       bool                          // kicked/parted
	NewGroupList GroupList                     // new GroupList
}

// GroupList, currently we detect spoofing by ensuring the origin of the
// message.  This may not be sufficient and we may have to add a signature of
// sorts.  For now roll with this assumption.
type GroupList struct {
	Name       string // group name
	Generation uint64 // incremented every time list changes
	Timestamp  int64  // unix time last generation changed

	// all participants, [0] is administrator
	// receiver must check [0] == originator
	Members [][zkidentity.IdentitySize]byte
}

// GroupMessage is a message to a group.
type GroupMessage struct {
	Name       string // group name
	Generation uint64 // Generation used
	Message    string // Actual message
	Mode       uint32 // 0 regular mode, 1 /me
}

// ChunkNew describes a chunked file transfer initiation.
type ChunkNew struct {
	Size        uint64            // total file size
	ChunkSize   uint64            // chunk size
	Filename    string            // original filename
	Description string            // user provided description
	MIME        string            // mime type
	Digest      [sha256.Size]byte // digest of file -> unique identifier
}

type Chunk struct {
	Offset  uint64            // offset in file
	Digest  [sha256.Size]byte // digest of file -> unique identifier
	Payload []byte            // chunk
}
<|endoftext|>


<commit_before>package rpc

import (
	"encoding/json"
	"errors"
	"fmt"
	. "github.com/zrepl/zrepl/util"
	"github.com/zrepl/zrepl/zfs"
	"io"
	"reflect"
)

type RPCRequester interface {
	FilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error)
	FilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error)
	InitialTransferRequest(r InitialTransferRequest) (io.Reader, error)
	IncrementalTransferRequest(r IncrementalTransferRequest) (io.Reader, error)
}

type RPCHandler interface {
	HandleFilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error)

	// returned versions ordered by birthtime, oldest first
	HandleFilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error)

	HandleInitialTransferRequest(r InitialTransferRequest) (io.Reader, error)
	HandleIncrementalTransferRequest(r IncrementalTransferRequest) (io.Reader, error)
}

const ByteStreamRPCProtocolVersion = 1

type ByteStreamRPC struct {
	conn    io.ReadWriteCloser
	encoder *json.Encoder
	decoder *json.Decoder
}

func ConnectByteStreamRPC(conn io.ReadWriteCloser) (RPCRequester, error) {
	// TODO do ssh connection to transport, establish TCP-like communication channel
	rpc := ByteStreamRPC{
		conn:    conn,
		encoder: json.NewEncoder(conn),
		decoder: json.NewDecoder(conn),
	}

	// Assert protocol versions are equal
	err := rpc.ProtocolVersionRequest()
	if err != nil {
		return nil, err
	}

	return rpc, nil
}

func ListenByteStreamRPC(conn io.ReadWriteCloser, handler RPCHandler) error {

	// A request consists of two subsequent JSON objects
	// Object 1: RequestHeader => contains type of Request Body
	// Object 2: RequestBody, e.g. IncrementalTransferRequest
	// A response is always a ResponseHeader followed by bytes to be interpreted
	// 	as indicated by the ResponseHeader.ResponseType, e.g.
	// 	 a) a chunked response
	//   b) or another JSON object

	defer conn.Close()

	decoder := json.NewDecoder(conn)
	encoder := json.NewEncoder(conn)

	for {

		var header RequestHeader = RequestHeader{}
		if err := decoder.Decode(&header); err != nil {
			respondWithError(encoder, EDecodeHeader, err)
			return conn.Close()
		}

		switch header.Type {
		case RTProtocolVersionRequest:
			var rq ByteStreamRPCProtocolVersionRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			if rq.ClientVersion != ByteStreamRPCProtocolVersion {
				respondWithError(encoder, EProtocolVersionMismatch, nil)
				return conn.Close()
			}

			r := ResponseHeader{
				RequestId:    header.Id,
				ResponseType: ROK,
			}
			if err := encoder.Encode(&r); err != nil {
				panic(err)
			}

		case RTFilesystemRequest:

			var rq FilesystemRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			roots, err := handler.HandleFilesystemRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
				return conn.Close()
			} else {
				r := ResponseHeader{
					RequestId:    header.Id,
					ResponseType: RFilesystems,
				}
				if err := encoder.Encode(&r); err != nil {
					panic(err)
				}
				if err := encoder.Encode(&roots); err != nil {
					panic(err)
				}
			}

		case RTFilesystemVersionsRequest:

			var rq FilesystemVersionsRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, err)
				return err
			}

			diff, err := handler.HandleFilesystemVersionsRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
				return err
			} else {
				r := ResponseHeader{
					RequestId:    header.Id,
					ResponseType: RFilesystemDiff,
				}
				if err := encoder.Encode(&r); err != nil {
					panic(err)
				}
				if err := encoder.Encode(&diff); err != nil {
					panic(err)
				}
			}

		case RTInitialTransferRequest:
			var rq InitialTransferRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			snapReader, err := handler.HandleInitialTransferRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
				return conn.Close()
			} else {
				chunker := NewChunker(snapReader)
				_, err := io.Copy(conn, &chunker)
				if err != nil {
					panic(err)
				}
			}

		case RTIncrementalTransferRequest:

			var rq IncrementalTransferRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			snapReader, err := handler.HandleIncrementalTransferRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
			} else {
				chunker := NewChunker(snapReader)
				_, err := io.Copy(conn, &chunker)
				if err != nil {
					panic(err)
				}
			}

		default:
			respondWithError(encoder, EUnknownRequestType, nil)
			return conn.Close()
		}
	}

	return nil
}

func respondWithError(encoder *json.Encoder, id ErrorId, err error) {

	r := ResponseHeader{
		ErrorId:      id,
		ResponseType: RNONE,
		Message:      err.Error(),
	}
	if err := encoder.Encode(&r); err != nil {
		panic(err)
	}

}

func inferRequestType(v interface{}) (RequestType, error) {
	switch v.(type) {
	case ByteStreamRPCProtocolVersionRequest:
		return RTProtocolVersionRequest, nil
	case FilesystemRequest:
		return RTFilesystemRequest, nil
	case FilesystemVersionsRequest:
		return RTFilesystemVersionsRequest, nil
	case InitialTransferRequest:
		return RTInitialTransferRequest, nil
	default:
		return 0, errors.New(fmt.Sprintf("cannot infer request type for type '%v'",
			reflect.TypeOf(v)))
	}
}

func genUUID() [16]byte {
	return [16]byte{} // TODO
}

func (c ByteStreamRPC) sendRequest(v interface{}) (err error) {

	var rt RequestType

	if rt, err = inferRequestType(v); err != nil {
		return
	}

	h := RequestHeader{
		Type: rt,
		Id:   genUUID(),
	}

	if err = c.encoder.Encode(h); err != nil {
		return
	}
	if err = c.encoder.Encode(v); err != nil {
		return
	}

	return
}

func (c ByteStreamRPC) expectResponseType(rt ResponseType) (err error) {
	var h ResponseHeader
	if err = c.decoder.Decode(&h); err != nil {
		return
	}

	if h.ResponseType != rt {
		return errors.New("unexpected response type in response header")
	}
	return
}

func (c ByteStreamRPC) sendRequestReceiveHeader(request interface{}, rt ResponseType) (err error) {

	if err = c.sendRequest(request); err != nil {
		return err
	}

	if err = c.expectResponseType(rt); err != nil {
		return err
	}

	return nil
}

func (c ByteStreamRPC) ProtocolVersionRequest() (err error) {
	b := ByteStreamRPCProtocolVersionRequest{
		ClientVersion: ByteStreamRPCProtocolVersion,
	}

	// OK response means the remote side can cope with our protocol version
	return c.sendRequestReceiveHeader(b, ROK)
}

func (c ByteStreamRPC) FilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error) {

	if err = c.sendRequestReceiveHeader(r, RFilesystems); err != nil {
		return
	}

	roots = make([]zfs.DatasetPath, 0)

	if err = c.decoder.Decode(&roots); err != nil {
		return
	}

	return
}

func (c ByteStreamRPC) FilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error) {

	if err = c.sendRequestReceiveHeader(r, RFilesystemDiff); err != nil {
		return
	}

	err = c.decoder.Decode(&versions)
	return
}

func (c ByteStreamRPC) InitialTransferRequest(r InitialTransferRequest) (unchunker io.Reader, err error) {

	if err = c.sendRequestReceiveHeader(r, RChunkedStream); err != nil {
		return
	}
	unchunker = NewUnchunker(c.conn)
	return
}

func (c ByteStreamRPC) IncrementalTransferRequest(r IncrementalTransferRequest) (unchunker io.Reader, err error) {
	if err = c.sendRequestReceiveHeader(r, RChunkedStream); err != nil {
		return
	}
	unchunker = NewUnchunker(c.conn)
	return
}

type LocalRPC struct {
	handler RPCHandler
}

func ConnectLocalRPC(handler RPCHandler) RPCRequester {
	return LocalRPC{handler}
}

func (c LocalRPC) FilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error) {
	return c.handler.HandleFilesystemRequest(r)
}

func (c LocalRPC) FilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error) {
	return c.handler.HandleFilesystemVersionsRequest(r)
}

func (c LocalRPC) InitialTransferRequest(r InitialTransferRequest) (io.Reader, error) {
	return c.handler.HandleInitialTransferRequest(r)
}

func (c LocalRPC) IncrementalTransferRequest(r IncrementalTransferRequest) (reader io.Reader, err error) {
	reader, err = c.handler.HandleIncrementalTransferRequest(r)
	return
}
<commit_msg>rpc: more detailed errors<commit_after>package rpc

import (
	"encoding/json"
	"errors"
	"fmt"
	. "github.com/zrepl/zrepl/util"
	"github.com/zrepl/zrepl/zfs"
	"io"
	"reflect"
)

type RPCRequester interface {
	FilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error)
	FilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error)
	InitialTransferRequest(r InitialTransferRequest) (io.Reader, error)
	IncrementalTransferRequest(r IncrementalTransferRequest) (io.Reader, error)
}

type RPCHandler interface {
	HandleFilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error)

	// returned versions ordered by birthtime, oldest first
	HandleFilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error)

	HandleInitialTransferRequest(r InitialTransferRequest) (io.Reader, error)
	HandleIncrementalTransferRequest(r IncrementalTransferRequest) (io.Reader, error)
}

const ByteStreamRPCProtocolVersion = 1

type ByteStreamRPC struct {
	conn    io.ReadWriteCloser
	encoder *json.Encoder
	decoder *json.Decoder
}

func ConnectByteStreamRPC(conn io.ReadWriteCloser) (RPCRequester, error) {
	// TODO do ssh connection to transport, establish TCP-like communication channel
	rpc := ByteStreamRPC{
		conn:    conn,
		encoder: json.NewEncoder(conn),
		decoder: json.NewDecoder(conn),
	}

	// Assert protocol versions are equal
	err := rpc.ProtocolVersionRequest()
	if err != nil {
		return nil, err
	}

	return rpc, nil
}

func ListenByteStreamRPC(conn io.ReadWriteCloser, handler RPCHandler) error {

	// A request consists of two subsequent JSON objects
	// Object 1: RequestHeader => contains type of Request Body
	// Object 2: RequestBody, e.g. IncrementalTransferRequest
	// A response is always a ResponseHeader followed by bytes to be interpreted
	// 	as indicated by the ResponseHeader.ResponseType, e.g.
	// 	 a) a chunked response
	//   b) or another JSON object

	defer conn.Close()

	decoder := json.NewDecoder(conn)
	encoder := json.NewEncoder(conn)

	for {

		var header RequestHeader = RequestHeader{}
		if err := decoder.Decode(&header); err != nil {
			respondWithError(encoder, EDecodeHeader, err)
			return conn.Close()
		}

		switch header.Type {
		case RTProtocolVersionRequest:
			var rq ByteStreamRPCProtocolVersionRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			if rq.ClientVersion != ByteStreamRPCProtocolVersion {
				respondWithError(encoder, EProtocolVersionMismatch, nil)
				return conn.Close()
			}

			r := ResponseHeader{
				RequestId:    header.Id,
				ResponseType: ROK,
			}
			if err := encoder.Encode(&r); err != nil {
				panic(err)
			}

		case RTFilesystemRequest:

			var rq FilesystemRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			roots, err := handler.HandleFilesystemRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
				return conn.Close()
			} else {
				r := ResponseHeader{
					RequestId:    header.Id,
					ResponseType: RFilesystems,
				}
				if err := encoder.Encode(&r); err != nil {
					panic(err)
				}
				if err := encoder.Encode(&roots); err != nil {
					panic(err)
				}
			}

		case RTFilesystemVersionsRequest:

			var rq FilesystemVersionsRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, err)
				return err
			}

			diff, err := handler.HandleFilesystemVersionsRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
				return err
			} else {
				r := ResponseHeader{
					RequestId:    header.Id,
					ResponseType: RFilesystemDiff,
				}
				if err := encoder.Encode(&r); err != nil {
					panic(err)
				}
				if err := encoder.Encode(&diff); err != nil {
					panic(err)
				}
			}

		case RTInitialTransferRequest:
			var rq InitialTransferRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			snapReader, err := handler.HandleInitialTransferRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
				return conn.Close()
			} else {
				chunker := NewChunker(snapReader)
				_, err := io.Copy(conn, &chunker)
				if err != nil {
					panic(err)
				}
			}

		case RTIncrementalTransferRequest:

			var rq IncrementalTransferRequest
			if err := decoder.Decode(&rq); err != nil {
				respondWithError(encoder, EDecodeRequestBody, nil)
				return conn.Close()
			}

			snapReader, err := handler.HandleIncrementalTransferRequest(rq)
			if err != nil {
				respondWithError(encoder, EHandler, err)
			} else {
				chunker := NewChunker(snapReader)
				_, err := io.Copy(conn, &chunker)
				if err != nil {
					panic(err)
				}
			}

		default:
			respondWithError(encoder, EUnknownRequestType, nil)
			return conn.Close()
		}
	}

	return nil
}

func respondWithError(encoder *json.Encoder, id ErrorId, err error) {

	r := ResponseHeader{
		ErrorId:      id,
		ResponseType: RNONE,
		Message:      err.Error(),
	}
	if err := encoder.Encode(&r); err != nil {
		panic(err)
	}

}

func inferRequestType(v interface{}) (RequestType, error) {
	switch v.(type) {
	case ByteStreamRPCProtocolVersionRequest:
		return RTProtocolVersionRequest, nil
	case FilesystemRequest:
		return RTFilesystemRequest, nil
	case FilesystemVersionsRequest:
		return RTFilesystemVersionsRequest, nil
	case InitialTransferRequest:
		return RTInitialTransferRequest, nil
	default:
		return 0, errors.New(fmt.Sprintf("cannot infer request type for type '%v'",
			reflect.TypeOf(v)))
	}
}

func genUUID() [16]byte {
	return [16]byte{} // TODO
}

func (c ByteStreamRPC) sendRequest(v interface{}) (err error) {

	var rt RequestType

	if rt, err = inferRequestType(v); err != nil {
		return
	}

	h := RequestHeader{
		Type: rt,
		Id:   genUUID(),
	}

	if err = c.encoder.Encode(h); err != nil {
		return
	}
	if err = c.encoder.Encode(v); err != nil {
		return
	}

	return
}

func (c ByteStreamRPC) expectResponseType(rt ResponseType) (err error) {
	var h ResponseHeader
	if err = c.decoder.Decode(&h); err != nil {
		return
	}

	if h.ResponseType != rt {
		return errors.New(fmt.Sprintf("unexpected response type in response header: got %#v, expected %#v. response header: %#v",
			h.ResponseType, rt, h))
	}
	return
}

func (c ByteStreamRPC) sendRequestReceiveHeader(request interface{}, rt ResponseType) (err error) {

	if err = c.sendRequest(request); err != nil {
		return err
	}

	if err = c.expectResponseType(rt); err != nil {
		return err
	}

	return nil
}

func (c ByteStreamRPC) ProtocolVersionRequest() (err error) {
	b := ByteStreamRPCProtocolVersionRequest{
		ClientVersion: ByteStreamRPCProtocolVersion,
	}

	// OK response means the remote side can cope with our protocol version
	return c.sendRequestReceiveHeader(b, ROK)
}

func (c ByteStreamRPC) FilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error) {

	if err = c.sendRequestReceiveHeader(r, RFilesystems); err != nil {
		return
	}

	roots = make([]zfs.DatasetPath, 0)

	if err = c.decoder.Decode(&roots); err != nil {
		return
	}

	return
}

func (c ByteStreamRPC) FilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error) {

	if err = c.sendRequestReceiveHeader(r, RFilesystemDiff); err != nil {
		return
	}

	err = c.decoder.Decode(&versions)
	return
}

func (c ByteStreamRPC) InitialTransferRequest(r InitialTransferRequest) (unchunker io.Reader, err error) {

	if err = c.sendRequestReceiveHeader(r, RChunkedStream); err != nil {
		return
	}
	unchunker = NewUnchunker(c.conn)
	return
}

func (c ByteStreamRPC) IncrementalTransferRequest(r IncrementalTransferRequest) (unchunker io.Reader, err error) {
	if err = c.sendRequestReceiveHeader(r, RChunkedStream); err != nil {
		return
	}
	unchunker = NewUnchunker(c.conn)
	return
}

type LocalRPC struct {
	handler RPCHandler
}

func ConnectLocalRPC(handler RPCHandler) RPCRequester {
	return LocalRPC{handler}
}

func (c LocalRPC) FilesystemRequest(r FilesystemRequest) (roots []zfs.DatasetPath, err error) {
	return c.handler.HandleFilesystemRequest(r)
}

func (c LocalRPC) FilesystemVersionsRequest(r FilesystemVersionsRequest) (versions []zfs.FilesystemVersion, err error) {
	return c.handler.HandleFilesystemVersionsRequest(r)
}

func (c LocalRPC) InitialTransferRequest(r InitialTransferRequest) (io.Reader, error) {
	return c.handler.HandleInitialTransferRequest(r)
}

func (c LocalRPC) IncrementalTransferRequest(r IncrementalTransferRequest) (reader io.Reader, err error) {
	reader, err = c.handler.HandleIncrementalTransferRequest(r)
	return
}
<|endoftext|>


<commit_before>package config
<commit_msg>Add basics.<commit_after>package config

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"path"
	"reflect"
	"unsafe"

	"pault.ag/go/debian/control"
)

//
func Load(name string, data interface{}) error {
	localUser, err := user.Current()
	if err != nil {
		return err
	}
	rcPath := path.Join(localUser.HomeDir, fmt.Sprintf(".%src", name))
	fd, err := os.Open(rcPath)
	if err != nil {
		return err
	}
	defer fd.Close()
	err = control.Unmarshal(data, fd)
	return err
}

func flagPointer(incoming reflect.Value, data *flag.FlagSet) error {
	if incoming.Type().Kind() == reflect.Ptr {
		return flagPointer(incoming.Elem(), data)
	}

	for i := 0; i < incoming.NumField(); i++ {
		field := incoming.Field(i)
		fieldType := incoming.Type().Field(i)

		if it := fieldType.Tag.Get("flag"); it != "" {
			/* Register the flag */
			switch field.Type().Kind() {
			case reflect.Int:
				data.IntVar(
					(*int)(unsafe.Pointer(field.Addr().Pointer())),
					it,
					int(field.Int()),
					fieldType.Tag.Get("description"),
				)
				continue
			case reflect.String:
				data.StringVar(
					(*string)(unsafe.Pointer(field.Addr().Pointer())),
					it,
					field.String(),
					fieldType.Tag.Get("description"),
				)
				continue
			default:
				return fmt.Errorf("Unknown type: %s", field.Type().Kind())
			}
		}
	}
	return nil
}

func Flag(data interface{}) (*flag.FlagSet, error) {
	dataValue := reflect.ValueOf(data)
	flagSet := flag.NewFlagSet(dataValue.Type().Name(), flag.ExitOnError)
	err := flagPointer(dataValue, flagSet)
	if err != nil {
		return nil, err
	}
	return flagSet, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"

	"github.com/hashicorp/go-multierror"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/efs"
	"github.com/aws/aws-sdk-go/service/elasticache"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/rds"
	"github.com/aws/aws-sdk-go/service/route53"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/sns"
	"github.com/aws/aws-sdk-go/service/sqs"
)

type Config struct {
	AccessKey  string
	SecretKey  string
	Token      string
	Region     string
	MaxRetries int

	AllowedAccountIds   []interface{}
	ForbiddenAccountIds []interface{}

	DynamoDBEndpoint string
}

type AWSClient struct {
	cloudwatchconn     *cloudwatch.CloudWatch
	cloudwatchlogsconn *cloudwatchlogs.CloudWatchLogs
	dynamodbconn       *dynamodb.DynamoDB
	ec2conn            *ec2.EC2
	ecsconn            *ecs.ECS
	efsconn            *efs.EFS
	elbconn            *elb.ELB
	autoscalingconn    *autoscaling.AutoScaling
	s3conn             *s3.S3
	sqsconn            *sqs.SQS
	snsconn            *sns.SNS
	r53conn            *route53.Route53
	region             string
	rdsconn            *rds.RDS
	iamconn            *iam.IAM
	kinesisconn        *kinesis.Kinesis
	elasticacheconn    *elasticache.ElastiCache
	lambdaconn         *lambda.Lambda
}

// Client configures and returns a fully initialized AWSClient
func (c *Config) Client() (interface{}, error) {
	var client AWSClient

	// Get the auth and region. This can fail if keys/regions were not
	// specified and we're attempting to use the environment.
	var errs []error

	log.Println("[INFO] Building AWS region structure")
	err := c.ValidateRegion()
	if err != nil {
		errs = append(errs, err)
	}

	if len(errs) == 0 {
		// store AWS region in client struct, for region specific operations such as
		// bucket storage in S3
		client.region = c.Region

		log.Println("[INFO] Building AWS auth structure")
		// We fetched all credential sources in Provider. If they are
		// available, they'll already be in c. See Provider definition.
		creds := credentials.NewStaticCredentials(c.AccessKey, c.SecretKey, c.Token)
		awsConfig := &aws.Config{
			Credentials: creds,
			Region:      aws.String(c.Region),
			MaxRetries:  aws.Int(c.MaxRetries),
		}

		log.Println("[INFO] Initializing IAM Connection")
		client.iamconn = iam.New(awsConfig)

		err := c.ValidateCredentials(client.iamconn)
		if err != nil {
			errs = append(errs, err)
		}

		awsDynamoDBConfig := &aws.Config{
			Credentials: creds,
			Region:      aws.String(c.Region),
			MaxRetries:  aws.Int(c.MaxRetries),
			Endpoint:    aws.String(c.DynamoDBEndpoint),
		}

		log.Println("[INFO] Initializing DynamoDB connection")
		client.dynamodbconn = dynamodb.New(awsDynamoDBConfig)

		log.Println("[INFO] Initializing ELB connection")
		client.elbconn = elb.New(awsConfig)

		log.Println("[INFO] Initializing S3 connection")
		client.s3conn = s3.New(awsConfig)

		log.Println("[INFO] Initializing SQS connection")
		client.sqsconn = sqs.New(awsConfig)

		log.Println("[INFO] Initializing SNS connection")
		client.snsconn = sns.New(awsConfig)

		log.Println("[INFO] Initializing RDS Connection")
		client.rdsconn = rds.New(awsConfig)

		log.Println("[INFO] Initializing Kinesis Connection")
		client.kinesisconn = kinesis.New(awsConfig)

		authErr := c.ValidateAccountId(client.iamconn)
		if authErr != nil {
			errs = append(errs, authErr)
		}

		log.Println("[INFO] Initializing AutoScaling connection")
		client.autoscalingconn = autoscaling.New(awsConfig)

		log.Println("[INFO] Initializing EC2 Connection")
		client.ec2conn = ec2.New(awsConfig)

		log.Println("[INFO] Initializing ECS Connection")
		client.ecsconn = ecs.New(awsConfig)

		log.Println("[INFO] Initializing EFS Connection")
		client.efsconn = efs.New(awsConfig)

		// aws-sdk-go uses v4 for signing requests, which requires all global
		// endpoints to use 'us-east-1'.
		// See http://docs.aws.amazon.com/general/latest/gr/sigv4_changes.html
		log.Println("[INFO] Initializing Route 53 connection")
		client.r53conn = route53.New(&aws.Config{
			Credentials: creds,
			Region:      aws.String("us-east-1"),
			MaxRetries:  aws.Int(c.MaxRetries),
		})

		log.Println("[INFO] Initializing Elasticache Connection")
		client.elasticacheconn = elasticache.New(awsConfig)

		log.Println("[INFO] Initializing Lambda Connection")
		client.lambdaconn = lambda.New(awsConfig)

		log.Println("[INFO] Initializing CloudWatch SDK connection")
		client.cloudwatchconn = cloudwatch.New(awsConfig)

		log.Println("[INFO] Initializing CloudWatch Logs connection")
		client.cloudwatchlogsconn = cloudwatchlogs.New(awsConfig)
	}

	if len(errs) > 0 {
		return nil, &multierror.Error{Errors: errs}
	}

	return &client, nil
}

// ValidateRegion returns an error if the configured region is not a
// valid aws region and nil otherwise.
func (c *Config) ValidateRegion() error {
	var regions = [11]string{"us-east-1", "us-west-2", "us-west-1", "eu-west-1",
		"eu-central-1", "ap-southeast-1", "ap-southeast-2", "ap-northeast-1",
		"sa-east-1", "cn-north-1", "us-gov-west-1"}

	for _, valid := range regions {
		if c.Region == valid {
			return nil
		}
	}
	return fmt.Errorf("Not a valid region: %s", c.Region)
}

// Validate credentials early and fail before we do any graph walking.
// In the case of an IAM role/profile with insuffecient privileges, fail
// silently
func (c *Config) ValidateCredentials(iamconn *iam.IAM) error {
	_, err := iamconn.GetUser(nil)

	if awsErr, ok := err.(awserr.Error); ok {

		if awsErr.Code() == "AccessDenied" || awsErr.Code() == "ValidationError" {
			log.Printf("[WARN] AccessDenied Error with iam.GetUser, assuming IAM profile")
			// User may be an IAM instance profile, or otherwise IAM role without the
			// GetUser permissions, so fail silently
			return nil
		}

		if awsErr.Code() == "SignatureDoesNotMatch" {
			return fmt.Errorf("Failed authenticating with AWS: please verify credentials")
		}
	}

	return err
}

// ValidateAccountId returns a context-specific error if the configured account
// id is explicitly forbidden or not authorised; and nil if it is authorised.
func (c *Config) ValidateAccountId(iamconn *iam.IAM) error {
	if c.AllowedAccountIds == nil && c.ForbiddenAccountIds == nil {
		return nil
	}

	log.Printf("[INFO] Validating account ID")

	out, err := iamconn.GetUser(nil)

	if err != nil {
		awsErr, _ := err.(awserr.Error)
		if awsErr.Code() == "ValidationError" {
			log.Printf("[WARN] ValidationError with iam.GetUser, assuming its an IAM profile")
			// User may be an IAM instance profile, so fail silently.
			// If it is an IAM instance profile
			// validating account might be superfluous
		} else {
			return fmt.Errorf("Failed getting account ID from IAM: %s", err)
			// return error if the account id is explicitly not authorised
		}
	}

	account_id := strings.Split(*out.User.Arn, ":")[4]

	if c.ForbiddenAccountIds != nil {
		for _, id := range c.ForbiddenAccountIds {
			if id == account_id {
				return fmt.Errorf("Forbidden account ID (%s)", id)
			}
		}
	}

	if c.AllowedAccountIds != nil {
		for _, id := range c.AllowedAccountIds {
			if id == account_id {
				return nil
			}
		}
		return fmt.Errorf("Account ID not allowed (%s)", account_id)
	}

	return nil
}
<commit_msg>AWS config forced to us-east-1 in variable, not inline.<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"

	"github.com/hashicorp/go-multierror"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/efs"
	"github.com/aws/aws-sdk-go/service/elasticache"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/rds"
	"github.com/aws/aws-sdk-go/service/route53"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/sns"
	"github.com/aws/aws-sdk-go/service/sqs"
)

type Config struct {
	AccessKey  string
	SecretKey  string
	Token      string
	Region     string
	MaxRetries int

	AllowedAccountIds   []interface{}
	ForbiddenAccountIds []interface{}

	DynamoDBEndpoint string
}

type AWSClient struct {
	cloudwatchconn     *cloudwatch.CloudWatch
	cloudwatchlogsconn *cloudwatchlogs.CloudWatchLogs
	dynamodbconn       *dynamodb.DynamoDB
	ec2conn            *ec2.EC2
	ecsconn            *ecs.ECS
	efsconn            *efs.EFS
	elbconn            *elb.ELB
	autoscalingconn    *autoscaling.AutoScaling
	s3conn             *s3.S3
	sqsconn            *sqs.SQS
	snsconn            *sns.SNS
	r53conn            *route53.Route53
	region             string
	rdsconn            *rds.RDS
	iamconn            *iam.IAM
	kinesisconn        *kinesis.Kinesis
	elasticacheconn    *elasticache.ElastiCache
	lambdaconn         *lambda.Lambda
}

// Client configures and returns a fully initialized AWSClient
func (c *Config) Client() (interface{}, error) {
	var client AWSClient

	// Get the auth and region. This can fail if keys/regions were not
	// specified and we're attempting to use the environment.
	var errs []error

	log.Println("[INFO] Building AWS region structure")
	err := c.ValidateRegion()
	if err != nil {
		errs = append(errs, err)
	}

	if len(errs) == 0 {
		// store AWS region in client struct, for region specific operations such as
		// bucket storage in S3
		client.region = c.Region

		log.Println("[INFO] Building AWS auth structure")
		// We fetched all credential sources in Provider. If they are
		// available, they'll already be in c. See Provider definition.
		creds := credentials.NewStaticCredentials(c.AccessKey, c.SecretKey, c.Token)
		awsConfig := &aws.Config{
			Credentials: creds,
			Region:      aws.String(c.Region),
			MaxRetries:  aws.Int(c.MaxRetries),
		}

		log.Println("[INFO] Initializing IAM Connection")
		client.iamconn = iam.New(awsConfig)

		err := c.ValidateCredentials(client.iamconn)
		if err != nil {
			errs = append(errs, err)
		}

		awsDynamoDBConfig := &aws.Config{
			Credentials: creds,
			Region:      aws.String(c.Region),
			MaxRetries:  aws.Int(c.MaxRetries),
			Endpoint:    aws.String(c.DynamoDBEndpoint),
		}
		// Some services exist only in us-east-1, e.g. because they manage
		// resources that can span across multiple regions, or because
		// signature format v4 requires region to be us-east-1 for global
		// endpoints:
		// http://docs.aws.amazon.com/general/latest/gr/sigv4_changes.html
		usEast1AwsConfig := &aws.Config{
			Credentials: creds,
			Region:      aws.String("us-east-1"),
			MaxRetries:  aws.Int(c.MaxRetries),
		}

		log.Println("[INFO] Initializing DynamoDB connection")
		client.dynamodbconn = dynamodb.New(awsDynamoDBConfig)

		log.Println("[INFO] Initializing ELB connection")
		client.elbconn = elb.New(awsConfig)

		log.Println("[INFO] Initializing S3 connection")
		client.s3conn = s3.New(awsConfig)

		log.Println("[INFO] Initializing SQS connection")
		client.sqsconn = sqs.New(awsConfig)

		log.Println("[INFO] Initializing SNS connection")
		client.snsconn = sns.New(awsConfig)

		log.Println("[INFO] Initializing RDS Connection")
		client.rdsconn = rds.New(awsConfig)

		log.Println("[INFO] Initializing Kinesis Connection")
		client.kinesisconn = kinesis.New(awsConfig)

		authErr := c.ValidateAccountId(client.iamconn)
		if authErr != nil {
			errs = append(errs, authErr)
		}

		log.Println("[INFO] Initializing AutoScaling connection")
		client.autoscalingconn = autoscaling.New(awsConfig)

		log.Println("[INFO] Initializing EC2 Connection")
		client.ec2conn = ec2.New(awsConfig)

		log.Println("[INFO] Initializing ECS Connection")
		client.ecsconn = ecs.New(awsConfig)

		log.Println("[INFO] Initializing EFS Connection")
		client.efsconn = efs.New(awsConfig)

		log.Println("[INFO] Initializing Route 53 connection")
		client.r53conn = route53.New(usEast1AwsConfig)

		log.Println("[INFO] Initializing Elasticache Connection")
		client.elasticacheconn = elasticache.New(awsConfig)

		log.Println("[INFO] Initializing Lambda Connection")
		client.lambdaconn = lambda.New(awsConfig)

		log.Println("[INFO] Initializing CloudWatch SDK connection")
		client.cloudwatchconn = cloudwatch.New(awsConfig)

		log.Println("[INFO] Initializing CloudWatch Logs connection")
		client.cloudwatchlogsconn = cloudwatchlogs.New(awsConfig)
	}

	if len(errs) > 0 {
		return nil, &multierror.Error{Errors: errs}
	}

	return &client, nil
}

// ValidateRegion returns an error if the configured region is not a
// valid aws region and nil otherwise.
func (c *Config) ValidateRegion() error {
	var regions = [11]string{"us-east-1", "us-west-2", "us-west-1", "eu-west-1",
		"eu-central-1", "ap-southeast-1", "ap-southeast-2", "ap-northeast-1",
		"sa-east-1", "cn-north-1", "us-gov-west-1"}

	for _, valid := range regions {
		if c.Region == valid {
			return nil
		}
	}
	return fmt.Errorf("Not a valid region: %s", c.Region)
}

// Validate credentials early and fail before we do any graph walking.
// In the case of an IAM role/profile with insuffecient privileges, fail
// silently
func (c *Config) ValidateCredentials(iamconn *iam.IAM) error {
	_, err := iamconn.GetUser(nil)

	if awsErr, ok := err.(awserr.Error); ok {

		if awsErr.Code() == "AccessDenied" || awsErr.Code() == "ValidationError" {
			log.Printf("[WARN] AccessDenied Error with iam.GetUser, assuming IAM profile")
			// User may be an IAM instance profile, or otherwise IAM role without the
			// GetUser permissions, so fail silently
			return nil
		}

		if awsErr.Code() == "SignatureDoesNotMatch" {
			return fmt.Errorf("Failed authenticating with AWS: please verify credentials")
		}
	}

	return err
}

// ValidateAccountId returns a context-specific error if the configured account
// id is explicitly forbidden or not authorised; and nil if it is authorised.
func (c *Config) ValidateAccountId(iamconn *iam.IAM) error {
	if c.AllowedAccountIds == nil && c.ForbiddenAccountIds == nil {
		return nil
	}

	log.Printf("[INFO] Validating account ID")

	out, err := iamconn.GetUser(nil)

	if err != nil {
		awsErr, _ := err.(awserr.Error)
		if awsErr.Code() == "ValidationError" {
			log.Printf("[WARN] ValidationError with iam.GetUser, assuming its an IAM profile")
			// User may be an IAM instance profile, so fail silently.
			// If it is an IAM instance profile
			// validating account might be superfluous
		} else {
			return fmt.Errorf("Failed getting account ID from IAM: %s", err)
			// return error if the account id is explicitly not authorised
		}
	}

	account_id := strings.Split(*out.User.Arn, ":")[4]

	if c.ForbiddenAccountIds != nil {
		for _, id := range c.ForbiddenAccountIds {
			if id == account_id {
				return fmt.Errorf("Forbidden account ID (%s)", id)
			}
		}
	}

	if c.AllowedAccountIds != nil {
		for _, id := range c.AllowedAccountIds {
			if id == account_id {
				return nil
			}
		}
		return fmt.Errorf("Account ID not allowed (%s)", account_id)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"os"
)

var VERSION = "v1.1.3"
var NAME = "git-hooks"
var TRIGGERS = [...]string{"applypatch-msg", "commit-msg", "post-applypatch", "post-checkout", "post-commit", "post-merge", "post-receive", "pre-applypatch", "pre-auto-gc", "pre-commit", "prepare-commit-msg", "pre-rebase", "pre-receive", "update", "pre-push"}

var CONTRIB_DIRNAME = "githooks-contrib"

var tplPreInstall = `#!/usr/bin/env bash
echo \"git hooks not installed in this repository.  Run 'git hooks --install' to install it or 'git hooks -h' for more information.\"`
var tplPostInstall = `#!/usr/bin/env bash
git-hooks run "$0" "$@"`

var ENV = os.Getenv("ENV")

var DIRS = map[string]string{
	"HomeTemplate":   ".git-template-with-git-hooks",
	"GlobalTemplate": "/usr/share/git-core/templates",
}

var GIT = map[string]string{
	"SetTemplateDir":    "config --global init.templatedir ",
	"GetTemplateDir":    "config --global --get init.templatedir",
	"UnsetTemplateDir":  "config --global --unset init.templatedir",
	"RemoveTemplateDir": "config --global --remove init",
	"FirstCommit":       "rev-list --max-parents=0 HEAD",
}

var MESSAGES = map[string]string{
	"NotGitRepo":     "Current directory is not a git repo",
	"Installed":      "Git hooks ARE installed in this repository.",
	"NotInstalled":   "Git hooks are NOT installed in this repository. (Run 'git hooks install' to install it)",
	"ExistHooks":     "hooks.old already exists, perhaps you already installed?",
	"NotExistHooks":  "Error, hooks.old doesn't exists, aborting uninstall to not destroy something",
	"Restore":        "Restore hooks.old",
	"SetTemplateDir": "Git global config init.templatedir is now set to ",
	"UpdateToDate":   "git-hooks is update to date",
	"Incompatible":   "Version backward incompatible, manually update required",
}

func isTestEnv() bool {
	return ENV == "test"
}
<commit_msg>bump version to 1.2.0<commit_after>package main

import (
	"os"
)

var VERSION = "v1.2.0"
var NAME = "git-hooks"
var TRIGGERS = [...]string{"applypatch-msg", "commit-msg", "post-applypatch", "post-checkout", "post-commit", "post-merge", "post-receive", "pre-applypatch", "pre-auto-gc", "pre-commit", "prepare-commit-msg", "pre-rebase", "pre-receive", "update", "pre-push"}

var CONTRIB_DIRNAME = "githooks-contrib"

var tplPreInstall = `#!/usr/bin/env bash
echo \"git hooks not installed in this repository.  Run 'git hooks --install' to install it or 'git hooks -h' for more information.\"`
var tplPostInstall = `#!/usr/bin/env bash
git-hooks run "$0" "$@"`

var ENV = os.Getenv("ENV")

var DIRS = map[string]string{
	"HomeTemplate":   ".git-template-with-git-hooks",
	"GlobalTemplate": "/usr/share/git-core/templates",
}

var GIT = map[string]string{
	"SetTemplateDir":    "config --global init.templatedir ",
	"GetTemplateDir":    "config --global --get init.templatedir",
	"UnsetTemplateDir":  "config --global --unset init.templatedir",
	"RemoveTemplateDir": "config --global --remove init",
	"FirstCommit":       "rev-list --max-parents=0 HEAD",
}

var MESSAGES = map[string]string{
	"NotGitRepo":     "Current directory is not a git repo",
	"Installed":      "Git hooks ARE installed in this repository.",
	"NotInstalled":   "Git hooks are NOT installed in this repository. (Run 'git hooks install' to install it)",
	"ExistHooks":     "hooks.old already exists, perhaps you already installed?",
	"NotExistHooks":  "Error, hooks.old doesn't exists, aborting uninstall to not destroy something",
	"Restore":        "Restore hooks.old",
	"SetTemplateDir": "Git global config init.templatedir is now set to ",
	"UpdateToDate":   "git-hooks is update to date",
	"Incompatible":   "Version backward incompatible, manually update required",
}

func isTestEnv() bool {
	return ENV == "test"
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mo

import (
	"fmt"
	"reflect"
	"regexp"

	"github.com/vmware/govmomi/vim25/methods"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

var arrayOfRegexp = regexp.MustCompile("ArrayOf(.*)$")

func anyTypeToValue(t interface{}) reflect.Value {
	rt := reflect.TypeOf(t)
	rv := reflect.ValueOf(t)

	// Dereference if ArrayOfXYZ type
	m := arrayOfRegexp.FindStringSubmatch(rt.Name())
	if len(m) > 0 {
		// ArrayOfXYZ type has single field named XYZ
		rv = rv.FieldByName(m[1])
		if !rv.IsValid() {
			panic(fmt.Sprintf("expected %s type to have field %s", m[0], m[1]))
		}
	}

	return rv
}

func buildValueMap(v reflect.Value, m map[string]reflect.Value) {
	t := v.Type().Elem()
	for i := 0; i < t.NumField(); i++ {
		sft := t.Field(i)

		// Recurse into embedded field
		if sft.Anonymous {
			buildValueMap(v.Elem().Field(i).Addr(), m)
			continue
		}

		tag := sft.Tag.Get("mo")
		if tag == "" {
			continue
		}

		m[tag] = v.Elem().Field(i)
	}
}

// Returns pointer to type t.
func objectContentToType(o types.ObjectContent) reflect.Value {
	t, ok := t[o.Obj.Type]
	if !ok {
		panic("unknown type: " + o.Obj.Type)
	}

	v := reflect.New(t)

	// Build map of property names to assignable reflect.Value
	m := make(map[string]reflect.Value)
	buildValueMap(v, m)

	for _, p := range o.PropSet {
		rv, ok := m[p.Name]
		if ok {
			pv := anyTypeToValue(p.Val)

			// If type is a pointer, create new instance of type
			if rv.Kind() == reflect.Ptr {
				rv.Set(reflect.New(rv.Type().Elem()))
				rv = rv.Elem()
			}

			// If type is an interface, check if pv implements it
			if rv.Kind() == reflect.Interface {
				rt := rv.Type()
				pt := pv.Type()
				if !pt.Implements(rt) {
					// Check if pointer to pv implements it
					if reflect.PtrTo(pt).Implements(rt) {
						npv := reflect.New(pt)
						npv.Elem().Set(pv)
						pv = npv
					} else {
						panic(fmt.Sprintf("type %s doesn't implement %s", pt.Name(), rt.Name()))
					}
				}
			}

			rv.Set(pv)
		}
	}

	return v
}

// RetrievePropertiesForRequest calls the RetrieveProperties method with the
// specified request and decodes the response struct into the value pointed to
// by dst.
func RetrievePropertiesForRequest(r soap.RoundTripper, req types.RetrieveProperties, dst interface{}) error {
	rt := reflect.TypeOf(dst)
	if rt.Kind() != reflect.Ptr {
		panic("need pointer")
	}

	rv := reflect.ValueOf(dst).Elem()
	if !rv.CanSet() {
		panic("cannot set dst")
	}

	isSlice := false
	switch rt.Elem().Kind() {
	case reflect.Struct:
	case reflect.Slice:
		isSlice = true
	default:
		panic("unexpected type")
	}

	res, err := methods.RetrieveProperties(r, &req)
	if err != nil {
		return err
	}

	if isSlice {
		for _, p := range res.Returnval {
			v := objectContentToType(p)
			rv.Set(reflect.Append(rv, v.Elem()))
		}
	} else {
		switch len(res.Returnval) {
		case 0:
		case 1:
			v := objectContentToType(res.Returnval[0])
			rv.Set(v.Elem())
		default:
			// If dst is not a slice, expect to receive 0 or 1 results
			panic("more than 1 result")
		}
	}

	return nil
}

// RetrieveProperties retrieves the properties of the managed object specified
// as obj and decodes the response struct into the value pointed to by dst.
func RetrieveProperties(r soap.RoundTripper, pc, obj types.ManagedObjectReference, dst interface{}) error {
	req := types.RetrieveProperties{
		This: pc,
		SpecSet: []types.PropertyFilterSpec{
			{
				ObjectSet: []types.ObjectSpec{
					{
						Obj:  obj,
						Skip: false,
					},
				},
				PropSet: []types.PropertySpec{
					{
						All:  true,
						Type: obj.Type,
					},
				},
			},
		},
	}

	return RetrievePropertiesForRequest(r, req, dst)
}
<commit_msg>Assign reference to self in managed objects<commit_after>/*
Copyright (c) 2014 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mo

import (
	"fmt"
	"reflect"
	"regexp"

	"github.com/vmware/govmomi/vim25/methods"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

var arrayOfRegexp = regexp.MustCompile("ArrayOf(.*)$")

func anyTypeToValue(t interface{}) reflect.Value {
	rt := reflect.TypeOf(t)
	rv := reflect.ValueOf(t)

	// Dereference if ArrayOfXYZ type
	m := arrayOfRegexp.FindStringSubmatch(rt.Name())
	if len(m) > 0 {
		// ArrayOfXYZ type has single field named XYZ
		rv = rv.FieldByName(m[1])
		if !rv.IsValid() {
			panic(fmt.Sprintf("expected %s type to have field %s", m[0], m[1]))
		}
	}

	return rv
}

func buildValueMap(v reflect.Value, m map[string]reflect.Value) {
	t := v.Type().Elem()
	for i := 0; i < t.NumField(); i++ {
		sft := t.Field(i)

		// Recurse into embedded field
		if sft.Anonymous {
			buildValueMap(v.Elem().Field(i).Addr(), m)
			continue
		}

		tag := sft.Tag.Get("mo")
		if tag == "" {
			continue
		}

		m[tag] = v.Elem().Field(i)
	}
}

// assignReference looks for the "ref" field in the specified struct and
// assigns the specified ManagedObjectReference.
//
// TODO(PN): buildValueMap and assignReference can be improved by combinding
// their functionality and having a type that only traversed the entire struct
// once. The type can then store the field indices (fields can be nested) for
// all the managed object references and the reference to itself.
//
func assignReference(v reflect.Value, ref types.ManagedObjectReference) bool {
	t := v.Type().Elem()
	for i := 0; i < t.NumField(); i++ {
		sft := t.Field(i)
		if sft.Anonymous {
			if assignReference(v.Elem().Field(i).Addr(), ref) {
				return true
			}
			continue
		}

		if sft.Name == "Ref" {
			v.Elem().Field(i).Set(reflect.ValueOf(ref))
			return true
		}
	}

	return false
}

// Returns pointer to type t.
func objectContentToType(o types.ObjectContent) reflect.Value {
	t, ok := t[o.Obj.Type]
	if !ok {
		panic("unknown type: " + o.Obj.Type)
	}

	v := reflect.New(t)

	// Assign reference to self
	assignReference(v, o.Obj)

	// Build map of property names to assignable reflect.Value
	m := make(map[string]reflect.Value)
	buildValueMap(v, m)

	for _, p := range o.PropSet {
		rv, ok := m[p.Name]
		if ok {
			pv := anyTypeToValue(p.Val)

			// If type is a pointer, create new instance of type
			if rv.Kind() == reflect.Ptr {
				rv.Set(reflect.New(rv.Type().Elem()))
				rv = rv.Elem()
			}

			// If type is an interface, check if pv implements it
			if rv.Kind() == reflect.Interface {
				rt := rv.Type()
				pt := pv.Type()
				if !pt.Implements(rt) {
					// Check if pointer to pv implements it
					if reflect.PtrTo(pt).Implements(rt) {
						npv := reflect.New(pt)
						npv.Elem().Set(pv)
						pv = npv
					} else {
						panic(fmt.Sprintf("type %s doesn't implement %s", pt.Name(), rt.Name()))
					}
				}
			}

			rv.Set(pv)
		}
	}

	return v
}

// RetrievePropertiesForRequest calls the RetrieveProperties method with the
// specified request and decodes the response struct into the value pointed to
// by dst.
func RetrievePropertiesForRequest(r soap.RoundTripper, req types.RetrieveProperties, dst interface{}) error {
	rt := reflect.TypeOf(dst)
	if rt.Kind() != reflect.Ptr {
		panic("need pointer")
	}

	rv := reflect.ValueOf(dst).Elem()
	if !rv.CanSet() {
		panic("cannot set dst")
	}

	isSlice := false
	switch rt.Elem().Kind() {
	case reflect.Struct:
	case reflect.Slice:
		isSlice = true
	default:
		panic("unexpected type")
	}

	res, err := methods.RetrieveProperties(r, &req)
	if err != nil {
		return err
	}

	if isSlice {
		for _, p := range res.Returnval {
			v := objectContentToType(p)
			rv.Set(reflect.Append(rv, v.Elem()))
		}
	} else {
		switch len(res.Returnval) {
		case 0:
		case 1:
			v := objectContentToType(res.Returnval[0])
			rv.Set(v.Elem())
		default:
			// If dst is not a slice, expect to receive 0 or 1 results
			panic("more than 1 result")
		}
	}

	return nil
}

// RetrieveProperties retrieves the properties of the managed object specified
// as obj and decodes the response struct into the value pointed to by dst.
func RetrieveProperties(r soap.RoundTripper, pc, obj types.ManagedObjectReference, dst interface{}) error {
	req := types.RetrieveProperties{
		This: pc,
		SpecSet: []types.PropertyFilterSpec{
			{
				ObjectSet: []types.ObjectSpec{
					{
						Obj:  obj,
						Skip: false,
					},
				},
				PropSet: []types.PropertySpec{
					{
						All:  true,
						Type: obj.Type,
					},
				},
			},
		},
	}

	return RetrievePropertiesForRequest(r, req, dst)
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package soap

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/vmware/govmomi/vim25/progress"
	"github.com/vmware/govmomi/vim25/types"
	"github.com/vmware/govmomi/vim25/xml"
)

type HasFault interface {
	Fault() *Fault
}

type RoundTripper interface {
	RoundTrip(req, res HasFault) error
}

type Client struct {
	http.Client

	u        url.URL
	insecure bool
	d        *debugContainer
}

func NewClient(u url.URL, insecure bool) *Client {
	c := Client{
		u:        u,
		insecure: insecure,
		d:        newDebug(),
	}

	if c.u.Scheme == "https" {
		c.Client.Transport = &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: c.insecure,
			},
		}
	}

	c.Jar, _ = cookiejar.New(nil)
	c.u.User = nil

	return &c
}

func (c *Client) URL() url.URL {
	return c.u
}

type marshaledClient struct {
	Cookies  []*http.Cookie
	URL      *url.URL
	Insecure bool
}

func (c *Client) MarshalJSON() ([]byte, error) {
	m := marshaledClient{
		Cookies:  c.Jar.Cookies(&c.u),
		URL:      &c.u,
		Insecure: c.insecure,
	}

	return json.Marshal(m)
}

func (c *Client) UnmarshalJSON(b []byte) error {
	var m marshaledClient

	err := json.Unmarshal(b, &m)
	if err != nil {
		return err
	}

	*c = *NewClient(*m.URL, m.Insecure)
	c.Jar.SetCookies(m.URL, m.Cookies)

	return nil
}

func (c *Client) RoundTrip(reqBody, resBody HasFault) error {
	var httpreq *http.Request
	var httpres *http.Response
	var err error

	reqEnv := Envelope{Body: reqBody}
	resEnv := Envelope{Body: resBody}

	// Create debugging context for this round trip
	d := c.d.newRoundTrip()
	if d.enabled() {
		defer d.done()
	}

	b, err := xml.Marshal(reqEnv)
	if err != nil {
		panic(err)
	}

	rawReqBody := io.MultiReader(strings.NewReader(xml.Header), bytes.NewReader(b))
	httpreq, err = http.NewRequest("POST", c.u.String(), rawReqBody)
	if err != nil {
		panic(err)
	}

	httpreq.Header.Set(`Content-Type`, `text/xml; charset="utf-8"`)
	httpreq.Header.Set(`SOAPAction`, `urn:vim25/5.5`)

	if d.enabled() {
		d.debugRequest(httpreq)
	}

	tstart := time.Now()
	httpres, err = c.Client.Do(httpreq)
	tstop := time.Now()

	if d.enabled() {
		d.logf("done in %6dms", tstop.Sub(tstart)/time.Millisecond)
	}

	if err != nil {
		return err
	}

	if d.enabled() {
		d.debugResponse(httpres)
	}

	// Close response regardless of what happens next
	defer httpres.Body.Close()

	dec := xml.NewDecoder(httpres.Body)
	dec.TypeFunc = types.TypeFunc()
	err = dec.Decode(&resEnv)
	if err != nil {
		return err
	}

	if f := resBody.Fault(); f != nil {
		return WrapSoapFault(f)
	}

	return err
}

// ParseURL wraps url.Parse to rewrite the URL.Host field
// In the case of VM guest uploads or NFC lease URLs, a Host
// field with a value of "*" is rewritten to the Client's URL.Host.
func (c *Client) ParseURL(urlStr string) (*url.URL, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	host := strings.Split(u.Host, ":")
	if host[0] == "*" {
		// Also use Client's port, to support port forwarding
		u.Host = c.URL().Host
	}

	return u, nil
}

type Upload struct {
	Type          string
	Method        string
	ContentLength int64
	Headers       map[string]string
	Progress      progress.Sinker
}

var DefaultUpload = Upload{
	Type:   "application/octet-stream",
	Method: "PUT",
}

// Upload PUTs the local file to the given URL
func (c *Client) Upload(f io.Reader, u *url.URL, param *Upload) error {
	var err error

	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, f, param.ContentLength)
		f = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	req, err := http.NewRequest(param.Method, u.String(), f)
	if err != nil {
		return err
	}

	req.ContentLength = param.ContentLength
	req.Header.Set("Content-Type", param.Type)

	for k, v := range param.Headers {
		req.Header.Add(k, v)
	}

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	switch res.StatusCode {
	case http.StatusOK:
	case http.StatusCreated:
	default:
		err = errors.New(res.Status)
	}

	return err
}

// UploadFile PUTs the local file to the given URL
func (c *Client) UploadFile(file string, u *url.URL, param *Upload) error {
	if param == nil {
		p := DefaultUpload // Copy since we set ContentLength
		param = &p
	}

	s, err := os.Stat(file)
	if err != nil {
		return err
	}

	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	param.ContentLength = s.Size()

	return c.Upload(f, u, param)
}

type Download struct {
	Method   string
	Progress progress.Sinker
}

var DefaultDownload = Download{
	Method: "GET",
}

// DownloadFile GETs the given URL to a local file
func (c *Client) DownloadFile(file string, u *url.URL, param *Download) error {
	var err error

	if param == nil {
		param = &DefaultDownload
	}

	fh, err := os.Create(file)
	if err != nil {
		return err
	}
	defer fh.Close()

	req, err := http.NewRequest(param.Method, u.String(), nil)
	if err != nil {
		return err
	}

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	defer res.Body.Close()

	switch res.StatusCode {
	case http.StatusOK:
	default:
		err = errors.New(res.Status)
	}

	if err != nil {
		return err
	}

	var r io.Reader = res.Body
	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, res.Body, res.ContentLength)
		r = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	_, err = io.Copy(fh, r)
	if err != nil {
		return err
	}

	// Assign error before returning so that it gets picked up by the deferred
	// function marking the progress reader as done.
	err = fh.Close()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Store reference to http.Transport<commit_after>/*
Copyright (c) 2014 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package soap

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/vmware/govmomi/vim25/progress"
	"github.com/vmware/govmomi/vim25/types"
	"github.com/vmware/govmomi/vim25/xml"
)

type HasFault interface {
	Fault() *Fault
}

type RoundTripper interface {
	RoundTrip(req, res HasFault) error
}

type Client struct {
	http.Client

	u url.URL
	k bool // Named after curl's -k flag
	d *debugContainer
	t *http.Transport
}

func NewClient(u url.URL, insecure bool) *Client {
	c := Client{
		u: u,
		k: insecure,
		d: newDebug(),
	}

	// Initialize http.RoundTripper on client, so we can customize it below
	c.t = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
	}

	if c.u.Scheme == "https" {
		c.t.TLSClientConfig = &tls.Config{InsecureSkipVerify: c.k}
		c.t.TLSHandshakeTimeout = 10 * time.Second
	}

	c.Client.Transport = c.t
	c.Client.Jar, _ = cookiejar.New(nil)
	c.u.User = nil

	return &c
}

func (c *Client) URL() url.URL {
	return c.u
}

type marshaledClient struct {
	Cookies  []*http.Cookie
	URL      *url.URL
	Insecure bool
}

func (c *Client) MarshalJSON() ([]byte, error) {
	m := marshaledClient{
		Cookies:  c.Jar.Cookies(&c.u),
		URL:      &c.u,
		Insecure: c.k,
	}

	return json.Marshal(m)
}

func (c *Client) UnmarshalJSON(b []byte) error {
	var m marshaledClient

	err := json.Unmarshal(b, &m)
	if err != nil {
		return err
	}

	*c = *NewClient(*m.URL, m.Insecure)
	c.Jar.SetCookies(m.URL, m.Cookies)

	return nil
}

func (c *Client) RoundTrip(reqBody, resBody HasFault) error {
	var httpreq *http.Request
	var httpres *http.Response
	var err error

	reqEnv := Envelope{Body: reqBody}
	resEnv := Envelope{Body: resBody}

	// Create debugging context for this round trip
	d := c.d.newRoundTrip()
	if d.enabled() {
		defer d.done()
	}

	b, err := xml.Marshal(reqEnv)
	if err != nil {
		panic(err)
	}

	rawReqBody := io.MultiReader(strings.NewReader(xml.Header), bytes.NewReader(b))
	httpreq, err = http.NewRequest("POST", c.u.String(), rawReqBody)
	if err != nil {
		panic(err)
	}

	httpreq.Header.Set(`Content-Type`, `text/xml; charset="utf-8"`)
	httpreq.Header.Set(`SOAPAction`, `urn:vim25/5.5`)

	if d.enabled() {
		d.debugRequest(httpreq)
	}

	tstart := time.Now()
	httpres, err = c.Client.Do(httpreq)
	tstop := time.Now()

	if d.enabled() {
		d.logf("done in %6dms", tstop.Sub(tstart)/time.Millisecond)
	}

	if err != nil {
		return err
	}

	if d.enabled() {
		d.debugResponse(httpres)
	}

	// Close response regardless of what happens next
	defer httpres.Body.Close()

	dec := xml.NewDecoder(httpres.Body)
	dec.TypeFunc = types.TypeFunc()
	err = dec.Decode(&resEnv)
	if err != nil {
		return err
	}

	if f := resBody.Fault(); f != nil {
		return WrapSoapFault(f)
	}

	return err
}

// ParseURL wraps url.Parse to rewrite the URL.Host field
// In the case of VM guest uploads or NFC lease URLs, a Host
// field with a value of "*" is rewritten to the Client's URL.Host.
func (c *Client) ParseURL(urlStr string) (*url.URL, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	host := strings.Split(u.Host, ":")
	if host[0] == "*" {
		// Also use Client's port, to support port forwarding
		u.Host = c.URL().Host
	}

	return u, nil
}

type Upload struct {
	Type          string
	Method        string
	ContentLength int64
	Headers       map[string]string
	Progress      progress.Sinker
}

var DefaultUpload = Upload{
	Type:   "application/octet-stream",
	Method: "PUT",
}

// Upload PUTs the local file to the given URL
func (c *Client) Upload(f io.Reader, u *url.URL, param *Upload) error {
	var err error

	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, f, param.ContentLength)
		f = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	req, err := http.NewRequest(param.Method, u.String(), f)
	if err != nil {
		return err
	}

	req.ContentLength = param.ContentLength
	req.Header.Set("Content-Type", param.Type)

	for k, v := range param.Headers {
		req.Header.Add(k, v)
	}

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	switch res.StatusCode {
	case http.StatusOK:
	case http.StatusCreated:
	default:
		err = errors.New(res.Status)
	}

	return err
}

// UploadFile PUTs the local file to the given URL
func (c *Client) UploadFile(file string, u *url.URL, param *Upload) error {
	if param == nil {
		p := DefaultUpload // Copy since we set ContentLength
		param = &p
	}

	s, err := os.Stat(file)
	if err != nil {
		return err
	}

	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	param.ContentLength = s.Size()

	return c.Upload(f, u, param)
}

type Download struct {
	Method   string
	Progress progress.Sinker
}

var DefaultDownload = Download{
	Method: "GET",
}

// DownloadFile GETs the given URL to a local file
func (c *Client) DownloadFile(file string, u *url.URL, param *Download) error {
	var err error

	if param == nil {
		param = &DefaultDownload
	}

	fh, err := os.Create(file)
	if err != nil {
		return err
	}
	defer fh.Close()

	req, err := http.NewRequest(param.Method, u.String(), nil)
	if err != nil {
		return err
	}

	res, err := c.Client.Do(req)
	if err != nil {
		return err
	}

	defer res.Body.Close()

	switch res.StatusCode {
	case http.StatusOK:
	default:
		err = errors.New(res.Status)
	}

	if err != nil {
		return err
	}

	var r io.Reader = res.Body
	if param.Progress != nil {
		pr := progress.NewReader(param.Progress, res.Body, res.ContentLength)
		r = pr

		// Mark progress reader as done when returning from this function.
		defer func() {
			pr.Done(err)
		}()
	}

	_, err = io.Copy(fh, r)
	if err != nil {
		return err
	}

	// Assign error before returning so that it gets picked up by the deferred
	// function marking the progress reader as done.
	err = fh.Close()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Package bridgeproxy relays with a remote TLS host via a HTTP proxy bridge.
package bridgeproxy

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
)

// Peer is a server we are connecting to. This can either be an
// intermediate http(s) proxy server or the final server we want
// to connect to.
type Peer struct {
	TLSConfig    *tls.Config // nil if unencrypted, valid config otherwise
	HostName     string      // The hostname to connect to
	Port         int         // The port to connect to on the hostname
	ConnectExtra string      // Extra headers to send after the CONNECT line
}

func forward(src io.ReadCloser, dst io.WriteCloser) {
	if _, err := io.Copy(dst, src); err != nil {
		log.Println("Could not forward:", err)
	}
	src.Close()
	dst.Close()
}

// readLine reads a network line one byte at a time. We need to read unbuffered
// as we might later turn the connection after a 200 response for a CONNECT
// into a tls connection, for which we need a net.Conn.
func readLine(src io.Reader) (string, error) {
	line := make([]byte, 0, 64)
	length := 0
	for length < 2 || line[length-2] != '\r' || line[length-1] != '\n' {
		line = append(line, 0)
		if _, err := io.ReadFull(src, line[length:]); err != nil {
			return "", err
		}
		length++
	}
	return string(line[:length]), nil
}

func writeResponse(w io.Writer, code int, format string, printf ...interface{}) {
	msg := fmt.Sprintf(format, printf...)
	fmt.Fprintf(w, "HTTP/1.0 %d %s\r\n", code, http.StatusText(code))
	fmt.Fprintf(w, "Content-Type: text/plain\r\n")
	fmt.Fprintf(w, "Content-Length: %d\r\n", len(msg))
	fmt.Fprintf(w, "\r\n")
	fmt.Fprintf(w, "%s", msg)
}

// DialProxy dials a proxy using the given slice of peers. It returns a
// network connection and error. Even if an error is returned, there may
// be a network connection that needs to be closed.
func DialProxy(peers []Peer) (net.Conn, error) {
	var connection net.Conn
	var err error
	for i, peer := range peers {
		// The first peer has to be dialed, others happen via connect
		if i == 0 {
			connection, err = net.Dial("tcp", fmt.Sprintf("%s:%d", peer.HostName, peer.Port))
			if err != nil {
				return connection, fmt.Errorf("could not dial proxy: %s\r\n", err)
			}
		} else {
			fmt.Fprintf(connection, "CONNECT %s:%d HTTP/1.0\r\n%s\r\n\r\n", peer.HostName, peer.Port, peers[i-1].ConnectExtra)

			line, err := readLine(io.LimitReader(connection, 1024))
			if err != nil {
				return connection, fmt.Errorf("could not CONNECT to %s: %s\r\n", peer.HostName, err.Error())
			}
			if !strings.HasPrefix(line, "HTTP/1.0 200") && !strings.HasPrefix(line, "HTTP/1.1 200") {
				return connection, fmt.Errorf("could not CONNECT to %s: %s", peer.HostName, line)
			}
			if _, err = readLine(connection); err != nil {
				return connection, fmt.Errorf("could not CONNECT to %s: Missing second line", peer.HostName)
			}
		}

		if peer.TLSConfig != nil {
			connection = tls.Client(connection, peer.TLSConfig)
		}
	}
	return connection, nil
}

// handleRequest handles a request by calling dialProxy() and then forwarding
func handleRequest(client net.Conn, peers []Peer) {
	remote, err := DialProxy(peers)
	if err != nil {
		log.Println("Error:", strings.TrimSpace(err.Error()))
		writeResponse(client, 502, "Error: %s", err.Error())
		if remote != nil {
			remote.Close()
		}
		client.Close()
		return
	}

	go forward(client, remote)
	go forward(remote, client)
}

// Serve serves the specified configuration, forwarding any packets from the
// local address given in listenAdress to the last peer specified in peers via
// any peers before specified before it.
func Serve(listenAdress string, peers []Peer) {
	// Listen for incoming connections.
	l, err := net.Listen("tcp", listenAdress)
	if err != nil {
		log.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	// Close the listener when the application closes.
	defer l.Close()
	log.Println("Listening on", listenAdress)
	log.Println("- Forwarding requests to", peers[len(peers)-1], "via", peers[0:len(peers)-1])
	for {
		// Listen for an incoming connection.
		conn, err := l.Accept()
		if err != nil {
			log.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		go handleRequest(conn, peers)
	}
}
<commit_msg>handle error when writing CONNECT line<commit_after>// Package bridgeproxy relays with a remote TLS host via a HTTP proxy bridge.
package bridgeproxy

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
)

// Peer is a server we are connecting to. This can either be an
// intermediate http(s) proxy server or the final server we want
// to connect to.
type Peer struct {
	TLSConfig    *tls.Config // nil if unencrypted, valid config otherwise
	HostName     string      // The hostname to connect to
	Port         int         // The port to connect to on the hostname
	ConnectExtra string      // Extra headers to send after the CONNECT line
}

func forward(src io.ReadCloser, dst io.WriteCloser) {
	if _, err := io.Copy(dst, src); err != nil {
		log.Println("Could not forward:", err)
	}
	src.Close()
	dst.Close()
}

// readLine reads a network line one byte at a time. We need to read unbuffered
// as we might later turn the connection after a 200 response for a CONNECT
// into a tls connection, for which we need a net.Conn.
func readLine(src io.Reader) (string, error) {
	line := make([]byte, 0, 64)
	length := 0
	for length < 2 || line[length-2] != '\r' || line[length-1] != '\n' {
		line = append(line, 0)
		if _, err := io.ReadFull(src, line[length:]); err != nil {
			return "", err
		}
		length++
	}
	return string(line[:length]), nil
}

func writeResponse(w io.Writer, code int, format string, printf ...interface{}) {
	msg := fmt.Sprintf(format, printf...)
	fmt.Fprintf(w, "HTTP/1.0 %d %s\r\n", code, http.StatusText(code))
	fmt.Fprintf(w, "Content-Type: text/plain\r\n")
	fmt.Fprintf(w, "Content-Length: %d\r\n", len(msg))
	fmt.Fprintf(w, "\r\n")
	fmt.Fprintf(w, "%s", msg)
}

// DialProxy dials a proxy using the given slice of peers. It returns a
// network connection and error. Even if an error is returned, there may
// be a network connection that needs to be closed.
func DialProxy(peers []Peer) (net.Conn, error) {
	var connection net.Conn
	var err error
	for i, peer := range peers {
		// The first peer has to be dialed, others happen via connect
		if i == 0 {
			connection, err = net.Dial("tcp", fmt.Sprintf("%s:%d", peer.HostName, peer.Port))
			if err != nil {
				return connection, fmt.Errorf("could not dial proxy: %s\r\n", err)
			}
		} else {
			if _, err := fmt.Fprintf(connection, "CONNECT %s:%d HTTP/1.0\r\n%s\r\n\r\n", peer.HostName, peer.Port, peers[i-1].ConnectExtra); err != nil {
				return connection, fmt.Errorf("failure writing CONNECT to %s: %s", peer.HostName, err.Error())
			}

			line, err := readLine(io.LimitReader(connection, 1024))
			if err != nil {
				return connection, fmt.Errorf("could not CONNECT to %s: %s\r\n", peer.HostName, err.Error())
			}
			if !strings.HasPrefix(line, "HTTP/1.0 200") && !strings.HasPrefix(line, "HTTP/1.1 200") {
				return connection, fmt.Errorf("could not CONNECT to %s: %s", peer.HostName, line)
			}
			if _, err = readLine(connection); err != nil {
				return connection, fmt.Errorf("could not CONNECT to %s: Missing second line", peer.HostName)
			}
		}

		if peer.TLSConfig != nil {
			connection = tls.Client(connection, peer.TLSConfig)
		}
	}
	return connection, nil
}

// handleRequest handles a request by calling dialProxy() and then forwarding
func handleRequest(client net.Conn, peers []Peer) {
	remote, err := DialProxy(peers)
	if err != nil {
		log.Println("Error:", strings.TrimSpace(err.Error()))
		writeResponse(client, 502, "Error: %s", err.Error())
		if remote != nil {
			remote.Close()
		}
		client.Close()
		return
	}

	go forward(client, remote)
	go forward(remote, client)
}

// Serve serves the specified configuration, forwarding any packets from the
// local address given in listenAdress to the last peer specified in peers via
// any peers before specified before it.
func Serve(listenAdress string, peers []Peer) {
	// Listen for incoming connections.
	l, err := net.Listen("tcp", listenAdress)
	if err != nil {
		log.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	// Close the listener when the application closes.
	defer l.Close()
	log.Println("Listening on", listenAdress)
	log.Println("- Forwarding requests to", peers[len(peers)-1], "via", peers[0:len(peers)-1])
	for {
		// Listen for an incoming connection.
		conn, err := l.Accept()
		if err != nil {
			log.Println("Error accepting: ", err.Error())
			os.Exit(1)
		}
		// Handle connections in a new goroutine.
		go handleRequest(conn, peers)
	}
}
<|endoftext|>