<commit_before>// Copyright (c) 2017 J. Hartsfield

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Package inboxer is a Go library for checking email using the google Gmail
// API.
package inboxer

// SCOPE:
// TODO:
// channels and go routines
// Mark as read/unread/important/spam
// check/return errors
//
//
// tests
// Watch inbox
// DOCS
// README.md
// how-to: add client credentials (for readme)
// Get Previews/snippet (put in docs)
//
// WORKS:
// Get emails by query
// Get email metadata
// Get email main body
// Get labels
// Check for unread messages
// Convert date to human readable format
//
// DONE:
// LICENSE

import (
	"encoding/base64"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	gmail "google.golang.org/api/gmail/v1"
)

// GetBody gets, decodes, and returns the body of the email. It returns an
// error if decoding goes wrong. mimeType is used to indicate whether you want
// the plain text or html encoding ("text/html", "text/plain").
func GetBody(msg *gmail.Message, mimeType string) (string, error) {
	for _, v := range msg.Payload.Parts {
		if v.MimeType == "multipart/alternative" {
			for _, l := range v.Parts {
				if l.MimeType == mimeType && l.Body.Size >= 1 {
					dec, err := decodeEmailBody(l.Body.Data)
					if err != nil {
						return "", err
					}
					return dec, nil
				}
			}
		}
		if v.MimeType == mimeType && v.Body.Size >= 1 {
			dec, err := decodeEmailBody(v.Body.Data)
			if err != nil {
				return "", err
			}
			return dec, nil
		}
	}
	return "", errors.New("Couldn't Read Body")
}

// HasLabel takes an email and a label and checks if that email has that label
func HasLabel(msg *gmail.Message, label string) bool {
	for _, v := range msg.LabelIds {
		if v == strings.ToUpper(label) {
			return true
		}
	}
	return false
}

// PartialMetadata stores email metadata. Some fields may sound redundant, but
// infact have different contexts. Some are slices of string because the ones
// that have multiple values are still being sorted from those that don't.
type PartialMetadata struct {
	// Sender is the entity that originally created and sent the message
	Sender string
	// From is the entitiy that sent the message to you (e.g. googlegroups). Most
	// of the time this information is only relevant to mailing lists.
	From string
	// Subject is the email subject
	Subject string
	// Mailing list contains the name of the mailing list that the email was
	// posted to, if any.
	MailingList string
	// CC is the "carbon copy" list of addresses
	CC []string
	// To is the recipient of the email.
	To []string
	// ThreadTopic contains the topic of the thread (e.g. google groups threads)
	ThreadTopic []string
	// DeliveredTo is who the email was sent to. This can contain multiple
	// addresses if the email was forwarded.
	DeliveredTo []string
}

// GetPartialMetadata gets some of the useful metadata from the headers.
func GetPartialMetadata(msg *gmail.Message) *PartialMetadata {
	info := &PartialMetadata{}
	for _, v := range msg.Payload.Headers {
		switch v.Name {
		case "Sender":
			info.Sender = v.Value
		case "From":
			info.From = v.Value
		case "Subject":
			info.Subject = v.Value
		case "Mailing-list":
			info.MailingList = v.Value
		case "CC":
			info.CC = append(info.CC, v.Value)
		case "To":
			info.To = append(info.To, v.Value)
		case "Thread-Topic":
			info.ThreadTopic = append(info.ThreadTopic, v.Value)
		case "Delivered-To":
			info.DeliveredTo = append(info.DeliveredTo, v.Value)
		}
	}
	return info
}

// decodeEmailBody is used to decode the email body by converting from
// URLEncoded base64 to a string.
func decodeEmailBody(data string) (string, error) {
	decoded, err := base64.URLEncoding.DecodeString(data)
	if err != nil {
		fmt.Println("decode error:", err)
		return "", err
	}
	return string(decoded), nil
}

// ReceivedTime parses and converts a unix time stamp into a human readable
// format ().
func ReceivedTime(datetime int64) time.Time {
	conv := strconv.FormatInt(datetime, 10)
	// Remove trailing zeros.
	conv = conv[:len(conv)-3]
	tc, err := strconv.ParseInt(conv, 10, 64)
	if err != nil {
		fmt.Println(err)
	}
	return time.Unix(tc, 0)
}

// Query queries the inbox for a string following the search style of the gmail
// online mailbox.
// example:
// "in:sent after:2017/01/01 before:2017/01/30"
func Query(srv *gmail.Service, query string) []*gmail.Message {
	inbox, err := srv.Users.Messages.List("me").Q(query).Do()
	if err != nil {
		fmt.Println(err)
	}
	return getById(srv, inbox)
}

// getById gets emails individually by ID. This is neccessary because this is
// how the gmail API is set up apparently (but why?).
//
func getById(srv *gmail.Service, msgs *gmail.ListMessagesResponse) []*gmail.Message {
	var msgSlice []*gmail.Message
	for _, v := range msgs.Messages {
		msg, _ := srv.Users.Messages.Get("me", v.Id).Do()
		msgSlice = append(msgSlice, msg)
	}
	return msgSlice
}

// GetMessages gets and returns gmail messages
func GetMessages(srv *gmail.Service, howMany uint) ([]*gmail.Message, error) {
	var msgSlice []*gmail.Message

	// Get the messages
	msgs, err := srv.Users.Messages.List("me").MaxResults(int64(howMany)).Do()
	if err != nil {
		return msgSlice, err
	}

	return getById(srv, msgs), nil
}

// CheckForUnreadByLabel checks for unread mail maching the speciified label.
// NOTE: When checking your inbox for unread messages, it's not uncommon for
// it return thousands of unread messages that you don't know about. To see them
// in gmail, search your mail for "label:unread". For CheckForUnreadByLabel to
// work properly you need to mark all mail as read either through gmail or
// through the MarkAllAsRead() function found in this library.
func CheckForUnreadByLabel(srv *gmail.Service, label string) (int64, error) {
	inbox, err := srv.Users.Labels.Get("me", label).Do()
	if err != nil {
		return -1, err
	}

	// fmt.Println(label.MessagesUnread)
	if inbox.MessagesUnread == 0 && inbox.ThreadsUnread == 0 {
		return 0, nil
	}

	return inbox.MessagesUnread + inbox.ThreadsUnread, nil
}

// CheckForUnread checks for mail labeled "UNREAD".
// NOTE: When checking your inbox for unread messages, it's not uncommon for
// it return thousands of unread messages that you don't know about. To see them
// in gmail, search your mail for "label:unread". For CheckForUnread to
// work properly you need to mark all mail as read either through gmail or
// through the MarkAllAsRead() function found in this library.
func CheckForUnread(srv *gmail.Service) (int64, error) {
	inbox, err := srv.Users.Labels.Get("me", "UNREAD").Do()
	if err != nil {
		return -1, err
	}

	// fmt.Println(label.MessagesUnread)
	if inbox.MessagesUnread == 0 && inbox.ThreadsUnread == 0 {
		return 0, nil
	}

	return inbox.MessagesUnread + inbox.ThreadsUnread, nil
}

// GetLabels gets a list of the labels used in the users inbox.
func GetLabels(srv *gmail.Service) (*gmail.ListLabelsResponse, error) {
	return srv.Users.Labels.List("me").Do()
}

// func watchInbox() {
// 	req := &gmail.WatchRequest{
// 		LabelFilterAction: "include",
// 		LabelIds:          []string{"UNREAD"},
// 		TopicName:         "gmailmsg",
// 	}

// 	wr, _ := srv.Users.Watch("me", req).Do()
// 	fmt.Println(wr.ForceSendFields)
// }
<commit_msg>adding sources to getById()<commit_after>// Copyright (c) 2017 J. Hartsfield

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Package inboxer is a Go library for checking email using the google Gmail
// API.
package inboxer

// SCOPE:
// TODO:
// channels and go routines
// Mark as read/unread/important/spam
// check/return errors
//
//
// tests
// Watch inbox
// DOCS
// README.md
// how-to: add client credentials (for readme)
// Get Previews/snippet (put in docs)
//
// WORKS:
// Get emails by query
// Get email metadata
// Get email main body
// Get labels
// Check for unread messages
// Convert date to human readable format
//
// DONE:
// LICENSE

import (
	"encoding/base64"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	gmail "google.golang.org/api/gmail/v1"
)

// GetBody gets, decodes, and returns the body of the email. It returns an
// error if decoding goes wrong. mimeType is used to indicate whether you want
// the plain text or html encoding ("text/html", "text/plain").
func GetBody(msg *gmail.Message, mimeType string) (string, error) {
	for _, v := range msg.Payload.Parts {
		if v.MimeType == "multipart/alternative" {
			for _, l := range v.Parts {
				if l.MimeType == mimeType && l.Body.Size >= 1 {
					dec, err := decodeEmailBody(l.Body.Data)
					if err != nil {
						return "", err
					}
					return dec, nil
				}
			}
		}
		if v.MimeType == mimeType && v.Body.Size >= 1 {
			dec, err := decodeEmailBody(v.Body.Data)
			if err != nil {
				return "", err
			}
			return dec, nil
		}
	}
	return "", errors.New("Couldn't Read Body")
}

// HasLabel takes an email and a label and checks if that email has that label
func HasLabel(msg *gmail.Message, label string) bool {
	for _, v := range msg.LabelIds {
		if v == strings.ToUpper(label) {
			return true
		}
	}
	return false
}

// PartialMetadata stores email metadata. Some fields may sound redundant, but
// infact have different contexts. Some are slices of string because the ones
// that have multiple values are still being sorted from those that don't.
type PartialMetadata struct {
	// Sender is the entity that originally created and sent the message
	Sender string
	// From is the entitiy that sent the message to you (e.g. googlegroups). Most
	// of the time this information is only relevant to mailing lists.
	From string
	// Subject is the email subject
	Subject string
	// Mailing list contains the name of the mailing list that the email was
	// posted to, if any.
	MailingList string
	// CC is the "carbon copy" list of addresses
	CC []string
	// To is the recipient of the email.
	To []string
	// ThreadTopic contains the topic of the thread (e.g. google groups threads)
	ThreadTopic []string
	// DeliveredTo is who the email was sent to. This can contain multiple
	// addresses if the email was forwarded.
	DeliveredTo []string
}

// GetPartialMetadata gets some of the useful metadata from the headers.
func GetPartialMetadata(msg *gmail.Message) *PartialMetadata {
	info := &PartialMetadata{}
	for _, v := range msg.Payload.Headers {
		switch v.Name {
		case "Sender":
			info.Sender = v.Value
		case "From":
			info.From = v.Value
		case "Subject":
			info.Subject = v.Value
		case "Mailing-list":
			info.MailingList = v.Value
		case "CC":
			info.CC = append(info.CC, v.Value)
		case "To":
			info.To = append(info.To, v.Value)
		case "Thread-Topic":
			info.ThreadTopic = append(info.ThreadTopic, v.Value)
		case "Delivered-To":
			info.DeliveredTo = append(info.DeliveredTo, v.Value)
		}
	}
	return info
}

// decodeEmailBody is used to decode the email body by converting from
// URLEncoded base64 to a string.
func decodeEmailBody(data string) (string, error) {
	decoded, err := base64.URLEncoding.DecodeString(data)
	if err != nil {
		fmt.Println("decode error:", err)
		return "", err
	}
	return string(decoded), nil
}

// ReceivedTime parses and converts a unix time stamp into a human readable
// format ().
func ReceivedTime(datetime int64) time.Time {
	conv := strconv.FormatInt(datetime, 10)
	// Remove trailing zeros.
	conv = conv[:len(conv)-3]
	tc, err := strconv.ParseInt(conv, 10, 64)
	if err != nil {
		fmt.Println(err)
	}
	return time.Unix(tc, 0)
}

// Query queries the inbox for a string following the search style of the gmail
// online mailbox.
// example:
// "in:sent after:2017/01/01 before:2017/01/30"
func Query(srv *gmail.Service, query string) []*gmail.Message {
	inbox, err := srv.Users.Messages.List("me").Q(query).Do()
	if err != nil {
		fmt.Println(err)
	}
	return getById(srv, inbox)
}

// getById gets emails individually by ID. This is neccessary because this is
// how the gmail API is set [0][1] up apparently (but why?).
// [0] https://developers.google.com/gmail/api/v1/reference/users/messages/get
// [1] https://stackoverflow.com/questions/36365172/message-payload-is-always-null-for-all-messages-how-do-i-get-this-data
func getById(srv *gmail.Service, msgs *gmail.ListMessagesResponse) []*gmail.Message {
	var msgSlice []*gmail.Message
	for _, v := range msgs.Messages {
		msg, _ := srv.Users.Messages.Get("me", v.Id).Do()
		msgSlice = append(msgSlice, msg)
	}
	return msgSlice
}

// GetMessages gets and returns gmail messages
func GetMessages(srv *gmail.Service, howMany uint) ([]*gmail.Message, error) {
	var msgSlice []*gmail.Message

	// Get the messages
	msgs, err := srv.Users.Messages.List("me").MaxResults(int64(howMany)).Do()
	if err != nil {
		return msgSlice, err
	}

	return getById(srv, msgs), nil
}

// CheckForUnreadByLabel checks for unread mail maching the speciified label.
// NOTE: When checking your inbox for unread messages, it's not uncommon for
// it return thousands of unread messages that you don't know about. To see them
// in gmail, search your mail for "label:unread". For CheckForUnreadByLabel to
// work properly you need to mark all mail as read either through gmail or
// through the MarkAllAsRead() function found in this library.
func CheckForUnreadByLabel(srv *gmail.Service, label string) (int64, error) {
	inbox, err := srv.Users.Labels.Get("me", label).Do()
	if err != nil {
		return -1, err
	}

	if inbox.MessagesUnread == 0 && inbox.ThreadsUnread == 0 {
		return 0, nil
	}

	return inbox.MessagesUnread + inbox.ThreadsUnread, nil
}

// CheckForUnread checks for mail labeled "UNREAD".
// NOTE: When checking your inbox for unread messages, it's not uncommon for
// it return thousands of unread messages that you don't know about. To see them
// in gmail, search your mail for "label:unread". For CheckForUnread to
// work properly you need to mark all mail as read either through gmail or
// through the MarkAllAsRead() function found in this library.
func CheckForUnread(srv *gmail.Service) (int64, error) {
	inbox, err := srv.Users.Labels.Get("me", "UNREAD").Do()
	if err != nil {
		return -1, err
	}

	if inbox.MessagesUnread == 0 && inbox.ThreadsUnread == 0 {
		return 0, nil
	}

	return inbox.MessagesUnread + inbox.ThreadsUnread, nil
}

// GetLabels gets a list of the labels used in the users inbox.
func GetLabels(srv *gmail.Service) (*gmail.ListLabelsResponse, error) {
	return srv.Users.Labels.List("me").Do()
}

// func watchInbox() {
// 	req := &gmail.WatchRequest{
// 		LabelFilterAction: "include",
// 		LabelIds:          []string{"UNREAD"},
// 		TopicName:         "gmailmsg",
// 	}

// 	wr, _ := srv.Users.Watch("me", req).Do()
// 	fmt.Println(wr.ForceSendFields)
// }
<|endoftext|>


<commit_before>package git

/*
extern int _govcs_gcrypt_init();
#cgo LDFLAGS: -lgcrypt
*/
import "C"
import (
	"fmt"
	"log"
	"os"
	"os/user"
	"strings"

	"crypto/md5"

	"golang.org/x/crypto/ssh"

	git2go "github.com/libgit2/git2go"
	"github.com/sourcegraph/go-vcs/vcs"
	"github.com/sourcegraph/go-vcs/vcs/gitcmd"
	sshutil "github.com/sourcegraph/go-vcs/vcs/ssh"
	"github.com/sourcegraph/go-vcs/vcs/util"
)

func init() {
	// Overwrite the git cloner to use the faster libgit2
	// implementation.
	vcs.RegisterCloner("git", func(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
		return Clone(url, dir, opt)
	})
}

func init() {
	// Initialize gcrypt for multithreaded operation. See
	// gcrypt_init.c for more information.
	rv := C._govcs_gcrypt_init()
	if rv != 0 {
		log.Fatal("gcrypt multithreaded init failed (see gcrypt_init.c)")
	}
}

func Clone(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
	clopt := git2go.CloneOptions{Bare: opt.Bare}

	rc, cfs, err := makeRemoteCallbacks(url, opt.RemoteOpts)
	if err != nil {
		return nil, err
	}
	if cfs != nil {
		defer cfs.run()
	}
	clopt.RemoteCallbacks = rc

	u, err := git2go.Clone(url, dir, &clopt)
	if err != nil {
		return nil, err
	}
	cr, err := gitcmd.Open(dir)
	if err != nil {
		return nil, err
	}
	r := &Repository{cr, u}
	if err := r.UpdateEverything(opt.RemoteOpts); err != nil {
		return nil, err
	}
	return r, nil
}

func (r *Repository) UpdateEverything(opt vcs.RemoteOpts) error {
	// TODO(sqs): allow use of a remote other than "origin"
	rm, err := r.u.LoadRemote("origin")
	if err != nil {
		return err
	}

	rc, cfs, err := makeRemoteCallbacks(rm.Url(), opt)
	if err != nil {
		return err
	}
	if cfs != nil {
		defer cfs.run()
	}
	if rc != nil {
		rm.SetCallbacks(rc)
	}

	if err := rm.Fetch([]string{"+refs/*:refs/*"}, nil, ""); err != nil {
		return err
	}

	return nil
}

type cleanupFuncs []func() error

func (f cleanupFuncs) run() error {
	for _, cf := range f {
		if err := cf(); err != nil {
			return err
		}
	}
	return nil
}

// makeRemoteCallbacks constructs the remote callbacks for libgit2
// remote operations. Currently the remote callbacks are trivial
// (empty) except when using an SSH remote.
//
// cleanupFuncs's run method should be called when the RemoteCallbacks
// struct is done being used. It is OK to ignore the error return.
func makeRemoteCallbacks(url string, opt vcs.RemoteOpts) (rc *git2go.RemoteCallbacks, cfs cleanupFuncs, err error) {
	defer func() {
		// Clean up if error; don't expect the caller to clean up if
		// we have a non-nil error.
		if err != nil {
			cfs.run()
		}
	}()

	if opt.SSH != nil {
		var privkeyFilename, pubkeyFilename string
		var privkeyFile, pubkeyFile *os.File
		var err error

		if opt.SSH.PrivateKey != nil {
			privkeyFilename, privkeyFile, err = util.WriteKeyTempFile(url, opt.SSH.PrivateKey)
			if err != nil {
				return nil, nil, err
			}
			cfs = append(cfs, privkeyFile.Close)
			cfs = append(cfs, func() error { return os.Remove(privkeyFile.Name()) })

			// Derive public key from private key if empty.
			if opt.SSH.PublicKey == nil {
				privKey, err := ssh.ParsePrivateKey(opt.SSH.PrivateKey)
				if err != nil {
					return nil, cfs, err
				}
				opt.SSH.PublicKey = ssh.MarshalAuthorizedKey(privKey.PublicKey())
			}

			pubkeyFilename, pubkeyFile, err = util.WriteKeyTempFile(url, opt.SSH.PublicKey)
			if err != nil {
				return nil, cfs, err
			}
			cfs = append(cfs, pubkeyFile.Close)
			cfs = append(cfs, func() error { return os.Remove(pubkeyFile.Name()) })
		}

		rc = &git2go.RemoteCallbacks{
			CredentialsCallback: func(url string, usernameFromURL string, allowedTypes git2go.CredType) (int, *git2go.Cred) {
				var username string
				if usernameFromURL != "" {
					username = usernameFromURL
				} else if opt.SSH.User != "" {
					username = opt.SSH.User
				} else {
					if username == "" {
						u, err := user.Current()
						if err == nil {
							username = u.Username
						}
					}
				}
				if allowedTypes&git2go.CredTypeSshKey != 0 && opt.SSH.PrivateKey != nil {
					rv, cred := git2go.NewCredSshKey(username, pubkeyFilename, privkeyFilename, "")
					return rv, &cred
				}
				log.Printf("No authentication available for git URL %q.", url)
				rv, cred := git2go.NewCredDefault()
				return rv, &cred
			},
			CertificateCheckCallback: func(cert *git2go.Certificate, valid bool, hostname string) int {
				// libgit2 currently always returns valid=false. It
				// may return valid=true in the future if it checks
				// host keys using known_hosts, but let's ignore valid
				// so we don't get that behavior unexpectedly.

				if InsecureSkipCheckVerifySSH {
					return 0
				}

				if cert == nil {
					return -1
				}

				if cert.Hostkey.Kind&git2go.HostkeyMD5 > 0 {
					keys, found := standardKnownHosts.Lookup(hostname)
					if found {
						hostFingerprint := md5String(cert.Hostkey.HashMD5)
						for _, key := range keys {
							knownFingerprint := md5String(md5.Sum(key.Marshal()))
							if hostFingerprint == knownFingerprint {
								return 0
							}
						}
					}
				}

				log.Printf("Invalid certificate for SSH host %s: %v.", hostname, cert)
				return -1
			},
		}
	}

	return rc, cfs, nil
}

// InsecureSkipCheckVerifySSH controls whether the client verifies the
// SSH server's certificate or host key. If InsecureSkipCheckVerifySSH
// is true, the program is susceptible to a man-in-the-middle
// attack. This should only be used for testing.
var InsecureSkipCheckVerifySSH bool

// standardKnownHosts contains known_hosts from the system known_hosts
// file and the user's known_hosts file.
var standardKnownHosts sshutil.KnownHosts

func init() {
	var err error
	standardKnownHosts, err = sshutil.ReadStandardKnownHostsFiles()
	if err != nil {
		log.Printf("Warning: failed to read standard SSH known_hosts files (%s). SSH host key checking will fail.", err)
	}
}

// md5String returns a formatted string representing the given md5Sum in hex
func md5String(md5Sum [16]byte) string {
	md5Str := fmt.Sprintf("% x", md5Sum)
	md5Str = strings.Replace(md5Str, " ", ":", -1)
	return md5Str
}
<commit_msg>update to use latest git2go<commit_after>package git

/*
extern int _govcs_gcrypt_init();
#cgo LDFLAGS: -lgcrypt
*/
import "C"
import (
	"fmt"
	"log"
	"os"
	"os/user"
	"strings"

	"crypto/md5"

	"golang.org/x/crypto/ssh"

	git2go "github.com/libgit2/git2go"
	"github.com/sourcegraph/go-vcs/vcs"
	"github.com/sourcegraph/go-vcs/vcs/gitcmd"
	sshutil "github.com/sourcegraph/go-vcs/vcs/ssh"
	"github.com/sourcegraph/go-vcs/vcs/util"
)

func init() {
	// Overwrite the git cloner to use the faster libgit2
	// implementation.
	vcs.RegisterCloner("git", func(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
		return Clone(url, dir, opt)
	})
}

func init() {
	// Initialize gcrypt for multithreaded operation. See
	// gcrypt_init.c for more information.
	rv := C._govcs_gcrypt_init()
	if rv != 0 {
		log.Fatal("gcrypt multithreaded init failed (see gcrypt_init.c)")
	}
}

func Clone(url, dir string, opt vcs.CloneOpt) (vcs.Repository, error) {
	clopt := git2go.CloneOptions{Bare: opt.Bare}

	rc, cfs, err := makeRemoteCallbacks(url, opt.RemoteOpts)
	if err != nil {
		return nil, err
	}
	if cfs != nil {
		defer cfs.run()
	}
	clopt.RemoteCallbacks = rc

	u, err := git2go.Clone(url, dir, &clopt)
	if err != nil {
		return nil, err
	}
	cr, err := gitcmd.Open(dir)
	if err != nil {
		return nil, err
	}
	r := &Repository{cr, u}
	if err := r.UpdateEverything(opt.RemoteOpts); err != nil {
		return nil, err
	}
	return r, nil
}

func (r *Repository) UpdateEverything(opt vcs.RemoteOpts) error {
	// TODO(sqs): allow use of a remote other than "origin"
	rm, err := r.u.LoadRemote("origin")
	if err != nil {
		return err
	}

	rc, cfs, err := makeRemoteCallbacks(rm.Url(), opt)
	if err != nil {
		return err
	}
	if cfs != nil {
		defer cfs.run()
	}
	if rc != nil {
		rm.SetCallbacks(rc)
	}

	if err := rm.Fetch([]string{"+refs/*:refs/*"}, nil, ""); err != nil {
		return err
	}

	return nil
}

type cleanupFuncs []func() error

func (f cleanupFuncs) run() error {
	for _, cf := range f {
		if err := cf(); err != nil {
			return err
		}
	}
	return nil
}

// makeRemoteCallbacks constructs the remote callbacks for libgit2
// remote operations. Currently the remote callbacks are trivial
// (empty) except when using an SSH remote.
//
// cleanupFuncs's run method should be called when the RemoteCallbacks
// struct is done being used. It is OK to ignore the error return.
func makeRemoteCallbacks(url string, opt vcs.RemoteOpts) (rc *git2go.RemoteCallbacks, cfs cleanupFuncs, err error) {
	defer func() {
		// Clean up if error; don't expect the caller to clean up if
		// we have a non-nil error.
		if err != nil {
			cfs.run()
		}
	}()

	if opt.SSH != nil {
		var privkeyFilename, pubkeyFilename string
		var privkeyFile, pubkeyFile *os.File
		var err error

		if opt.SSH.PrivateKey != nil {
			privkeyFilename, privkeyFile, err = util.WriteKeyTempFile(url, opt.SSH.PrivateKey)
			if err != nil {
				return nil, nil, err
			}
			cfs = append(cfs, privkeyFile.Close)
			cfs = append(cfs, func() error { return os.Remove(privkeyFile.Name()) })

			// Derive public key from private key if empty.
			if opt.SSH.PublicKey == nil {
				privKey, err := ssh.ParsePrivateKey(opt.SSH.PrivateKey)
				if err != nil {
					return nil, cfs, err
				}
				opt.SSH.PublicKey = ssh.MarshalAuthorizedKey(privKey.PublicKey())
			}

			pubkeyFilename, pubkeyFile, err = util.WriteKeyTempFile(url, opt.SSH.PublicKey)
			if err != nil {
				return nil, cfs, err
			}
			cfs = append(cfs, pubkeyFile.Close)
			cfs = append(cfs, func() error { return os.Remove(pubkeyFile.Name()) })
		}

		rc = &git2go.RemoteCallbacks{
			CredentialsCallback: git2go.CredentialsCallback(func(url string, usernameFromURL string, allowedTypes git2go.CredType) (git2go.ErrorCode, *git2go.Cred) {
				var username string
				if usernameFromURL != "" {
					username = usernameFromURL
				} else if opt.SSH.User != "" {
					username = opt.SSH.User
				} else {
					if username == "" {
						u, err := user.Current()
						if err == nil {
							username = u.Username
						}
					}
				}
				if allowedTypes&git2go.CredTypeSshKey != 0 && opt.SSH.PrivateKey != nil {
					rv, cred := git2go.NewCredSshKey(username, pubkeyFilename, privkeyFilename, "")
					return git2go.ErrorCode(rv), &cred
				}
				log.Printf("No authentication available for git URL %q.", url)
				rv, cred := git2go.NewCredDefault()
				return git2go.ErrorCode(rv), &cred
			}),
			CertificateCheckCallback: git2go.CertificateCheckCallback(func(cert *git2go.Certificate, valid bool, hostname string) git2go.ErrorCode {
				// libgit2 currently always returns valid=false. It
				// may return valid=true in the future if it checks
				// host keys using known_hosts, but let's ignore valid
				// so we don't get that behavior unexpectedly.

				if InsecureSkipCheckVerifySSH {
					return git2go.ErrOk
				}

				if cert == nil {
					return git2go.ErrNotFound
				}

				if cert.Hostkey.Kind&git2go.HostkeyMD5 > 0 {
					keys, found := standardKnownHosts.Lookup(hostname)
					if found {
						hostFingerprint := md5String(cert.Hostkey.HashMD5)
						for _, key := range keys {
							knownFingerprint := md5String(md5.Sum(key.Marshal()))
							if hostFingerprint == knownFingerprint {
								return git2go.ErrOk
							}
						}
					}
				}

				log.Printf("Invalid certificate for SSH host %s: %v.", hostname, cert)
				return git2go.ErrGeneric
			}),
		}
	}

	return rc, cfs, nil
}

// InsecureSkipCheckVerifySSH controls whether the client verifies the
// SSH server's certificate or host key. If InsecureSkipCheckVerifySSH
// is true, the program is susceptible to a man-in-the-middle
// attack. This should only be used for testing.
var InsecureSkipCheckVerifySSH bool

// standardKnownHosts contains known_hosts from the system known_hosts
// file and the user's known_hosts file.
var standardKnownHosts sshutil.KnownHosts

func init() {
	var err error
	standardKnownHosts, err = sshutil.ReadStandardKnownHostsFiles()
	if err != nil {
		log.Printf("Warning: failed to read standard SSH known_hosts files (%s). SSH host key checking will fail.", err)
	}
}

// md5String returns a formatted string representing the given md5Sum in hex
func md5String(md5Sum [16]byte) string {
	md5Str := fmt.Sprintf("% x", md5Sum)
	md5Str = strings.Replace(md5Str, " ", ":", -1)
	return md5Str
}
<|endoftext|>


<commit_before>package zeroconf

import (
	"fmt"
	"log"
	"net"
	"strings"

	dns "github.com/miekg/godns"
)

var (
	Registry = newRegistry()
)

func openIPv4Socket() *net.UDPConn {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{
		IP:   net.IPv4zero,
		Port: 5353,
	})
	if err != nil {
		log.Fatal(err)
	}
	return conn
}

func init() {
	go Registry.mainloop()

	socket := openIPv4Socket()
	if err := socket.JoinGroup(nil, net.IPv4(224, 0, 0, 251)); err != nil {
		log.Fatal(err)
	}

	listener := &listener{
		socket:   socket,
		registry: Registry,
	}
	go listener.mainloop()

	// simple logger
	go func() {
		for op := range Registry.Subscribe() {
			switch op.Op {
			case Add:
				log.Printf("Add: %#v", op.Service)
			case Remove:
				log.Printf("Remove: %#v", op.Service)
			}
		}
	}()
}

func mcastInterfaces() []net.Interface {
	ifaces := make([]net.Interface, 0)
	interfaces, err := net.Interfaces()
	if err != nil {
		log.Fatal(err)
	}
	for _, i := range interfaces {
		if isMulticast(i) {
			fmt.Printf("%#v\n", i)
			ifaces = append(ifaces, i)
		}
	}
	return ifaces
}

func isMulticast(i net.Interface) bool {
	return (i.Flags&net.FlagUp > 0) && (i.Flags&net.FlagMulticast > 0)
}

type op int

const (
	Add    op = 1
	Remove op = 2
)

type Operation struct {
	Op      op
	Service *Service
}

type registry struct {
	ops         chan Operation
	subscribe   chan *subscription
	services    []*Service
	hosts	 []*Host
	subscribers []*subscription
}

// TODO registry is not an exported type, should it be?
func newRegistry() *registry {
	return &registry{
		ops:         make(chan Operation),
		subscribe:   make(chan *subscription),
		services:    nil,
		hosts:	nil,
		subscribers: nil,
	}
}

func (r *registry) AddService(service *Service) {
	r.ops <- Operation{
		Op:      Add,
		Service: service,
	}
}

func (r *registry) RemoveService(service *Service) {
	r.ops <- Operation{
		Op:      Remove,
		Service: service,
	}
}

// TODO subscribe should take a *Query
func (r *registry) Subscribe() chan Operation {
	s := &subscription{
		make(chan Operation),
	}
	r.subscribe <- s
	return s.c
}

func (r *registry) mainloop() {
	for {
		select {
		case op := <-r.ops:
			switch op.Op {
			case Add:
				r.addService(op.Service)
			}
			r.notifySubscribers(op)
		case sub := <-r.subscribe:
			r.subscribers = append(r.subscribers, sub)
		}
	}
}

func (r *registry) addService(service *Service) {
	r.services = append(r.services, service)
}

func (r *registry) notifySubscribers(op Operation) {
	for i, _ := range r.subscribers {
		r.subscribers[i].notify(op)
	}
}

type subscription struct {
	c chan Operation
}

func (s *subscription) notify(op Operation) {
	s.c <- op // TODO use non blocking send in case reciver is full
}

type listener struct {
	socket   *net.UDPConn
	registry *registry
}

func (l *listener) mainloop() {
	buf := make([]byte, 1500)
	for {
		read, err := l.socket.Read(buf)
		if err != nil {
			log.Fatal(err)
		}
		msg := new(dns.Msg)
		msg.Unpack(buf[:read])
		s := new(Service)
		s.unmarshal(msg)
		if s.valid() {
			l.registry.AddService(s)
		}
	}
}

type Service struct {
	Type       string
	Name       string
	Domain     string
	Host       string
	Port		uint16
	Additional []string
}

// s.unmarshal may not be complete, return false if so
func (s *Service) valid() bool {
	return len(s.Name) > 0
}

func (s *Service) unmarshal(msg *dns.Msg) {
	for i := range msg.Answer {
		switch rr := msg.Answer[i].(type) {
		case *dns.RR_SRV:
			x := strings.Split(rr.Hdr.Name, ".")
			s.Name = x[0]
			s.Type = strings.Join(x[1:3], ".")
			s.Domain = strings.Join(x[3:], ".")
			s.Host = rr.Target
			s.Port = rr.Port

		case *dns.RR_TXT:
			s.Additional = append(s.Additional, strings.Split(rr.Txt, ",")...)
		default:
			log.Printf("%#v", rr)
		}
	}
}

type Host struct {
	Addrs []*net.Addr
	Name string
}
<commit_msg>Simplified the subscription type<commit_after>package zeroconf

import (
	"fmt"
	"log"
	"net"
	"strings"

	dns "github.com/miekg/godns"
)

var (
	Registry = newRegistry()
)

func openIPv4Socket() *net.UDPConn {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{
		IP:   net.IPv4zero,
		Port: 5353,
	})
	if err != nil {
		log.Fatal(err)
	}
	return conn
}

func init() {
	go Registry.mainloop()

	socket := openIPv4Socket()
	if err := socket.JoinGroup(nil, net.IPv4(224, 0, 0, 251)); err != nil {
		log.Fatal(err)
	}

	listener := &listener{
		socket:   socket,
		registry: Registry,
	}
	go listener.mainloop()

	// simple logger
	go func() {
		for op := range Registry.Subscribe() {
			switch op.Op {
			case Add:
				log.Printf("Add: %#v", op.Service)
			case Remove:
				log.Printf("Remove: %#v", op.Service)
			}
		}
	}()
}

func mcastInterfaces() []net.Interface {
	ifaces := make([]net.Interface, 0)
	interfaces, err := net.Interfaces()
	if err != nil {
		log.Fatal(err)
	}
	for _, i := range interfaces {
		if isMulticast(i) {
			fmt.Printf("%#v\n", i)
			ifaces = append(ifaces, i)
		}
	}
	return ifaces
}

func isMulticast(i net.Interface) bool {
	return (i.Flags&net.FlagUp > 0) && (i.Flags&net.FlagMulticast > 0)
}

type op int

const (
	Add    op = 1
	Remove op = 2
)

type Operation struct {
	Op      op
	Service *Service
}

type registry struct {
	ops         chan Operation
	subscribe   chan subscription
	services    []*Service
	hosts       []*Host
	subscribers []subscription
}

// TODO registry is not an exported type, should it be?
func newRegistry() *registry {
	return &registry{
		ops:         make(chan Operation),
		subscribe:   make(chan subscription),
		services:    nil,
		hosts:       nil,
		subscribers: nil,
	}
}

func (r *registry) AddService(service *Service) {
	r.ops <- Operation{
		Op:      Add,
		Service: service,
	}
}

func (r *registry) RemoveService(service *Service) {
	r.ops <- Operation{
		Op:      Remove,
		Service: service,
	}
}

// TODO subscribe should take a *Query
func (r *registry) Subscribe() chan Operation {
	s := make(chan Operation)
	r.subscribe <- s
	return s
}

func (r *registry) mainloop() {
	for {
		select {
		case op := <-r.ops:
			switch op.Op {
			case Add:
				r.addService(op.Service)
			}
			r.notifySubscribers(op)
		case sub := <-r.subscribe:
			r.subscribers = append(r.subscribers, sub)
		}
	}
}

func (r *registry) addService(service *Service) {
	r.services = append(r.services, service)
}

func (r *registry) notifySubscribers(op Operation) {
	for i, _ := range r.subscribers {
		//TODO use non blocking send in case reciever is full
		r.subscribers[i] <- op
	}
}

type subscription chan Operation

type listener struct {
	socket   *net.UDPConn
	registry *registry
}

func (l *listener) mainloop() {
	buf := make([]byte, 1500)
	for {
		read, err := l.socket.Read(buf)
		if err != nil {
			log.Fatal(err)
		}
		msg := new(dns.Msg)
		msg.Unpack(buf[:read])
		s := new(Service)
		s.unmarshal(msg)
		if s.valid() {
			l.registry.AddService(s)
		}
	}
}

type Service struct {
	Type       string
	Name       string
	Domain     string
	Host       string
	Port       uint16
	Additional []string
}

// s.unmarshal may not be complete, return false if so
func (s *Service) valid() bool {
	return len(s.Name) > 0
}

func (s *Service) unmarshal(msg *dns.Msg) {
	for i := range msg.Answer {
		switch rr := msg.Answer[i].(type) {
		case *dns.RR_SRV:
			x := strings.Split(rr.Hdr.Name, ".")
			s.Name = x[0]
			s.Type = strings.Join(x[1:3], ".")
			s.Domain = strings.Join(x[3:], ".")
			s.Host = rr.Target
			s.Port = rr.Port

		case *dns.RR_TXT:
			s.Additional = append(s.Additional, strings.Split(rr.Txt, ",")...)
		default:
			log.Printf("%#v", rr)
		}
	}
}

type Host struct {
	Addrs []*net.Addr
	Name  string
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors 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 args has common command-line flags for generation programs.
package args

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"k8s.io/kubernetes/cmd/libs/go2idl/generator"
	"k8s.io/kubernetes/cmd/libs/go2idl/namer"
	"k8s.io/kubernetes/cmd/libs/go2idl/parser"

	"github.com/spf13/pflag"
)

// Default returns a defaulted GeneratorArgs. You may change the defaults
// before calling AddFlags.
func Default() *GeneratorArgs {
	return &GeneratorArgs{
		OutputBase:       DefaultSourceTree(),
		GoHeaderFilePath: filepath.Join(DefaultSourceTree(), "k8s.io/kubernetes/hack/boilerplate/boilerplate.go.txt"),
	}
}

// GeneratorArgs has arguments common to most generators.
type GeneratorArgs struct {
	// Which directories to parse.
	InputDirs []string

	// Source tree to write results to.
	OutputBase string

	// Package path within the source tree.
	OutputPackagePath string

	// Where to get copyright header text.
	GoHeaderFilePath string

	// If true, only verify, don't write anything.
	VerifyOnly bool
}

func (g *GeneratorArgs) AddFlags(fs *pflag.FlagSet) {
	fs.StringSliceVarP(&g.InputDirs, "input-dirs", "i", g.InputDirs, "Comma-separated list of import paths to get input types from.")
	fs.StringVarP(&g.OutputBase, "output-base", "o", g.OutputBase, "Output base; defaults to $GOPATH/src/ or ./ if $GOPATH is not set.")
	fs.StringVarP(&g.OutputPackagePath, "output-package", "p", g.OutputPackagePath, "Base package path.")
	fs.StringVarP(&g.GoHeaderFilePath, "go-header-file", "h", g.GoHeaderFilePath, "File containing boilerplate header text. The string YEAR will be replaced with the current 4-digit year.")
	fs.BoolVar(&g.VerifyOnly, "verify-only", g.VerifyOnly, "If true, only verify existing output, do not write anything.")
}

// LoadGoBoilerplate loads the boilerplate file passed to --go-header-file.
func (g *GeneratorArgs) LoadGoBoilerplate() ([]byte, error) {
	b, err := ioutil.ReadFile(g.GoHeaderFilePath)
	if err != nil {
		return nil, err
	}
	b = bytes.Replace(b, []byte("YEAR"), []byte(strconv.Itoa(time.Now().Year())), -1)
	return b, nil
}

// NewBuilder makes a new parser.Builder and populates it with the input
// directories.
func (g *GeneratorArgs) NewBuilder() (*parser.Builder, error) {
	b := parser.New()
	for _, d := range g.InputDirs {
		if err := b.AddDir(d); err != nil {
			return nil, fmt.Errorf("unable to add directory %q: %v", d, err)
		}
	}
	return b, nil
}

// DefaultSourceTree returns the /src directory of the first entry in $GOPATH.
// If $GOPATH is empty, it returns "./". Useful as a default output location.
func DefaultSourceTree() string {
	paths := strings.Split(os.Getenv("GOPATH"), string(filepath.ListSeparator))
	if len(paths) > 0 && len(paths[0]) > 0 {
		return filepath.Join(paths[0], "src")
	}
	return "./"
}

// Execute implements main().
// If you don't need any non-default behavior, use as:
// args.Default().Execute(...)
func (g *GeneratorArgs) Execute(nameSystems namer.NameSystems, defaultSystem string, pkgs func(*generator.Context, *GeneratorArgs) generator.Packages) error {
	g.AddFlags(pflag.CommandLine)
	pflag.Parse()

	b, err := g.NewBuilder()
	if err != nil {
		return fmt.Errorf("Failed making a parser: %v", err)
	}

	c, err := generator.NewContext(b, nameSystems, defaultSystem)
	if err != nil {
		return fmt.Errorf("Failed making a context: %v", err)
	}

	c.Verify = g.VerifyOnly
	packages := pkgs(c, g)
	if err := c.ExecutePackages(g.OutputBase, packages); err != nil {
		return fmt.Errorf("Failed executing generator: %v", err)
	}

	return nil
}
<commit_msg>Update/verify client generator in scripts<commit_after>/*
Copyright 2015 The Kubernetes Authors 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 args has common command-line flags for generation programs.
package args

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"k8s.io/kubernetes/cmd/libs/go2idl/generator"
	"k8s.io/kubernetes/cmd/libs/go2idl/namer"
	"k8s.io/kubernetes/cmd/libs/go2idl/parser"

	"github.com/spf13/pflag"
)

// Default returns a defaulted GeneratorArgs. You may change the defaults
// before calling AddFlags.
func Default() *GeneratorArgs {
	generatorArgs := &GeneratorArgs{
		OutputBase:       DefaultSourceTree(),
		GoHeaderFilePath: filepath.Join(DefaultSourceTree(), "k8s.io/kubernetes/hack/boilerplate/boilerplate.go.txt"),
	}
	generatorArgs.AddFlags(pflag.CommandLine)
	return generatorArgs
}

// GeneratorArgs has arguments common to most generators.
type GeneratorArgs struct {
	// Which directories to parse.
	InputDirs []string

	// Source tree to write results to.
	OutputBase string

	// Package path within the source tree.
	OutputPackagePath string

	// Where to get copyright header text.
	GoHeaderFilePath string

	// If true, only verify, don't write anything.
	VerifyOnly bool
}

func (g *GeneratorArgs) AddFlags(fs *pflag.FlagSet) {
	fs.StringSliceVarP(&g.InputDirs, "input-dirs", "i", g.InputDirs, "Comma-separated list of import paths to get input types from.")
	fs.StringVarP(&g.OutputBase, "output-base", "o", g.OutputBase, "Output base; defaults to $GOPATH/src/ or ./ if $GOPATH is not set.")
	fs.StringVarP(&g.OutputPackagePath, "output-package", "p", g.OutputPackagePath, "Base package path.")
	fs.StringVarP(&g.GoHeaderFilePath, "go-header-file", "h", g.GoHeaderFilePath, "File containing boilerplate header text. The string YEAR will be replaced with the current 4-digit year.")
	fs.BoolVar(&g.VerifyOnly, "verify-only", g.VerifyOnly, "If true, only verify existing output, do not write anything.")
}

// LoadGoBoilerplate loads the boilerplate file passed to --go-header-file.
func (g *GeneratorArgs) LoadGoBoilerplate() ([]byte, error) {
	b, err := ioutil.ReadFile(g.GoHeaderFilePath)
	if err != nil {
		return nil, err
	}
	b = bytes.Replace(b, []byte("YEAR"), []byte(strconv.Itoa(time.Now().Year())), -1)
	return b, nil
}

// NewBuilder makes a new parser.Builder and populates it with the input
// directories.
func (g *GeneratorArgs) NewBuilder() (*parser.Builder, error) {
	b := parser.New()
	for _, d := range g.InputDirs {
		if err := b.AddDir(d); err != nil {
			return nil, fmt.Errorf("unable to add directory %q: %v", d, err)
		}
	}
	return b, nil
}

// DefaultSourceTree returns the /src directory of the first entry in $GOPATH.
// If $GOPATH is empty, it returns "./". Useful as a default output location.
func DefaultSourceTree() string {
	paths := strings.Split(os.Getenv("GOPATH"), string(filepath.ListSeparator))
	if len(paths) > 0 && len(paths[0]) > 0 {
		return filepath.Join(paths[0], "src")
	}
	return "./"
}

// Execute implements main().
// If you don't need any non-default behavior, use as:
// args.Default().Execute(...)
func (g *GeneratorArgs) Execute(nameSystems namer.NameSystems, defaultSystem string, pkgs func(*generator.Context, *GeneratorArgs) generator.Packages) error {
	pflag.Parse()

	b, err := g.NewBuilder()
	if err != nil {
		return fmt.Errorf("Failed making a parser: %v", err)
	}

	c, err := generator.NewContext(b, nameSystems, defaultSystem)
	if err != nil {
		return fmt.Errorf("Failed making a context: %v", err)
	}

	c.Verify = g.VerifyOnly
	packages := pkgs(c, g)
	if err := c.ExecutePackages(g.OutputBase, packages); err != nil {
		return fmt.Errorf("Failed executing generator: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package idpot

import (
  "crypto/sha1"
  "fmt"
  "io"
  "io/ioutil"
  "math/rand"
  "os"
  "os/exec"
  "runtime"
  "strings"
  "sync"
  "syscall"
  "testing"
  "time"
  "github.com/lestrrat/go-tcptest"
  "github.com/lestrrat/go-test-mysqld"
  "github.com/lestrrat/go-idpot/client"
)

func init () {
  rand.Seed(time.Now().UnixNano())
  buildCommandLineTools()

  // Create a mysqld instance for testing
}

func randomName() string {
  h := sha1.New()
  io.WriteString(h,
    fmt.Sprintf(
      "idpot.test.%d.%d",
      os.Getpid(),
      time.Now().UnixNano(),
    ),
  )

  return fmt.Sprintf("%x", h.Sum(nil))
}

func buildCommandLineTools() {
  // compile command
  cmd := exec.Command("go", "build", "-o", "bin/idpot-server", "cli/idpot-server.go")
  err := cmd.Run()
  if err != nil {
    panic(fmt.Sprintf("Failed to compile idpot-server: %s", err))
  }

  // Must add it to PATH so we can invoke it
  os.Setenv("PATH",
    strings.Join([]string {
      "bin",
      os.Getenv("PATH"),
    }, ":",),
  )
}

func TestServer(t *testing.T) {
  mysqld, err := mysqltest.NewMysqld(nil)
  if err != nil {
    t.Fatalf("Failed to start mysqld: %s", err)
  }
  defer mysqld.Stop()

  fullpath, err := exec.LookPath("idpot-server")
  if err != nil {
    t.Skipf("Could not find idpot-server in PATH")
  }

  // Create a config file to be used by idpot-server
  dir, err := ioutil.TempDir("", "idpot-test")
  if err != nil {
    t.Fatalf("Failed to create temporary directory: %s", err)
  }
  defer os.RemoveAll(dir)

  cfgfile, err := ioutil.TempFile(dir, "idpot.gcfg")
  if err != nil {
    t.Fatalf("Failed to create temp file: %s", err)
  }

  io.WriteString(cfgfile, fmt.Sprintf(
    `
[Server]
LogFile=idpot-%%Y%%m%%d-access_log
LogLinkName=idpot-access_log

[Mysql]
ConnectString=%s
`,
    mysqld.Datasource("test", "", "", 0),
  ))
  cfgfile.Close()

  var cmd *exec.Cmd
  start := func(port int) {
    numprocs := runtime.GOMAXPROCS(-1)
    os.Setenv("GOMAXPROCS", fmt.Sprintf("%d", numprocs))
    cmd = exec.Command(
      fullpath,
      "--listen", fmt.Sprintf("127.0.0.1:%d", port),
      "--config", cfgfile.Name(),
    )
    t.Logf("Starting command %v", cmd.Args)
    cmd.SysProcAttr = &syscall.SysProcAttr {
      Setpgid: true,
    }
    out, _ := cmd.CombinedOutput()
    t.Logf(string(out))
  }

  // Launch server in separate process
  server, err := tcptest.Start(start, 30 * time.Second)
  if err != nil {
    t.Fatalf("Failed to start idpot-server: %s", err)
  }

  port := server.Port()

  // First, declare a new pot
  base   := fmt.Sprintf("http://127.0.0.1:%d", port)
  client := client.New(base)

  for _, min := range []uint64 { 0, 100, 1000, 10000 } {
    pot    := randomName()
    err     = client.CreatePot(pot, min)
    if err != nil {
      t.Fatalf("Failed to create new pot %s: %s", pot, err)
    }

    t.Logf("Created new pot %s", pot)

    sem := make(chan bool, 1)

    wg := &sync.WaitGroup {}
    ids := make(map[uint64]bool)
    for i := 0; i < 10; i++ {
      wg.Add(1)
      go func() {
        defer func() { wg.Done() }()
        for j := 0; j < 100; j++ {
          id, err := client.NextId(pot)
          if err != nil {
            t.Errorf("Failed to fetch id: %s", err)
            continue
          }

          sem<- true
          if _, ok := ids[id]; ok {
            // collision!
            t.Errorf("Collision detected for id %d :/", id)
          } else {
            ids[id] = true
          }
          <-sem
        }
      }()
    }

    wg.Wait()

    // I should have 1000 ids
    count := 0
    for v, _ := range ids {
      if v < min {
        t.Errorf("Id given is out of range: %d (wanted > %d)", v, min)
      }
      count++
    }

    if count != 1000 {
      t.Errorf("Did not get 1000 ids")
    }
  }

  cmd.Process.Signal(syscall.SIGTERM)
  server.Wait()
}<commit_msg>show fetches/sec<commit_after>package idpot

import (
  "crypto/sha1"
  "fmt"
  "io"
  "io/ioutil"
  "math/rand"
  "os"
  "os/exec"
  "runtime"
  "strings"
  "sync"
  "syscall"
  "testing"
  "time"
  "github.com/lestrrat/go-tcptest"
  "github.com/lestrrat/go-test-mysqld"
  "github.com/lestrrat/go-idpot/client"
)

func init () {
  rand.Seed(time.Now().UnixNano())
  buildCommandLineTools()

  // Create a mysqld instance for testing
}

func randomName() string {
  h := sha1.New()
  io.WriteString(h,
    fmt.Sprintf(
      "idpot.test.%d.%d",
      os.Getpid(),
      time.Now().UnixNano(),
    ),
  )

  return fmt.Sprintf("%x", h.Sum(nil))
}

func buildCommandLineTools() {
  // compile command
  cmd := exec.Command("go", "build", "-o", "bin/idpot-server", "cli/idpot-server.go")
  err := cmd.Run()
  if err != nil {
    panic(fmt.Sprintf("Failed to compile idpot-server: %s", err))
  }

  // Must add it to PATH so we can invoke it
  os.Setenv("PATH",
    strings.Join([]string {
      "bin",
      os.Getenv("PATH"),
    }, ":",),
  )
}

func TestServer(t *testing.T) {
  mysqld, err := mysqltest.NewMysqld(nil)
  if err != nil {
    t.Fatalf("Failed to start mysqld: %s", err)
  }
  defer mysqld.Stop()

  fullpath, err := exec.LookPath("idpot-server")
  if err != nil {
    t.Skipf("Could not find idpot-server in PATH")
  }

  // Create a config file to be used by idpot-server
  dir, err := ioutil.TempDir("", "idpot-test")
  if err != nil {
    t.Fatalf("Failed to create temporary directory: %s", err)
  }
  defer os.RemoveAll(dir)

  cfgfile, err := ioutil.TempFile(dir, "idpot.gcfg")
  if err != nil {
    t.Fatalf("Failed to create temp file: %s", err)
  }

  io.WriteString(cfgfile, fmt.Sprintf(
    `
[Server]
LogFile=idpot-%%Y%%m%%d-access_log
LogLinkName=idpot-access_log

[Mysql]
ConnectString=%s
`,
    mysqld.Datasource("test", "", "", 0),
  ))
  cfgfile.Close()

  var cmd *exec.Cmd
  start := func(port int) {
    numprocs := runtime.GOMAXPROCS(-1)
    os.Setenv("GOMAXPROCS", fmt.Sprintf("%d", numprocs))
    cmd = exec.Command(
      fullpath,
      "--listen", fmt.Sprintf("127.0.0.1:%d", port),
      "--config", cfgfile.Name(),
    )
    t.Logf("Starting command %v", cmd.Args)
    cmd.SysProcAttr = &syscall.SysProcAttr {
      Setpgid: true,
    }
    out, _ := cmd.CombinedOutput()
    t.Logf(string(out))
  }

  // Launch server in separate process
  server, err := tcptest.Start(start, 30 * time.Second)
  if err != nil {
    t.Fatalf("Failed to start idpot-server: %s", err)
  }

  port := server.Port()

  // First, declare a new pot
  base   := fmt.Sprintf("http://127.0.0.1:%d", port)
  client := client.New(base)

  for _, min := range []uint64 { 0, 100, 1000, 10000 } {
    pot    := randomName()
    err     = client.CreatePot(pot, min)
    if err != nil {
      t.Fatalf("Failed to create new pot %s: %s", pot, err)
    }

    t.Logf("Created new pot %s", pot)

    sem := make(chan bool, 1)

    wg := &sync.WaitGroup {}
    ids := make(map[uint64]bool)

    maxGoros := 10
    maxFetches := 10000

    t0 := time.Now()
    for i := 0; i < maxGoros; i++ {
      wg.Add(1)
      go func() {
        defer func() { wg.Done() }()
        for j := 0; j < maxFetches; j++ {
          id, err := client.NextId(pot)
          if err != nil {
            t.Errorf("Failed to fetch id: %s", err)
            continue
          }

          sem<- true
          if _, ok := ids[id]; ok {
            // collision!
            t.Errorf("Collision detected for id %d :/", id)
          } else {
            ids[id] = true
          }
          <-sem
        }
      }()
    }

    wg.Wait()
    elapsed := time.Since(t0)

    // I should have 1000 ids
    count := 0
    for v, _ := range ids {
      if v < min {
        t.Errorf("Id given is out of range: %d (wanted > %d)", v, min)
      }
      count++
    }

    if count != maxGoros * maxFetches {
      t.Errorf("Did not get %d ids, got %d", maxGoros * maxFetches, count)
    }

    t.Logf("Fetched %d ids in %f secs (%f fetches/sec)", count, elapsed.Seconds(), float64(count) / elapsed.Seconds())
  }

  cmd.Process.Signal(syscall.SIGTERM)
  server.Wait()
}<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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 adapter

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	scan "knative.dev/eventing-redis/source/pkg/redis"

	cloudevents "github.com/cloudevents/sdk-go/v2"
	"github.com/gomodule/redigo/redis"
	"go.uber.org/zap"
	"knative.dev/eventing/pkg/adapter/v2"
	"knative.dev/pkg/logging"
)

const (
	// RedisStreamSourceEventType is the default RedisStreamSource CloudEvent type.
	RedisStreamSourceEventType = "dev.knative.sources.redisstream"
)

func NewEnvConfig() adapter.EnvConfigAccessor {
	return &Config{}
}

type Adapter struct {
	config *Config
	logger *zap.Logger
	client cloudevents.Client
	source string
}

func NewAdapter(ctx context.Context, processed adapter.EnvConfigAccessor, ceClient cloudevents.Client) adapter.Adapter {
	config := processed.(*Config)
	return &Adapter{
		config: config,
		logger: logging.FromContext(ctx).Desugar().With(zap.String("stream", config.Stream)),
		client: ceClient,
		source: fmt.Sprintf("%s/%s", config.Address, config.Stream),
	}
}

func (a *Adapter) Start(ctx context.Context) error {
	pool := newPool(a.config.Address)

	conn, err := pool.Dial()
	if err != nil {
		return err
	}

	// Build consumer group name from stateful set pod name of adapter
	groupName := a.config.PodName

	a.logger.Info("Retrieving group info", zap.String("group", groupName))
	groups, err := scan.ScanXInfoGroupReply(conn.Do("XINFO", "GROUPS", a.config.Stream))

	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "no such key") || strings.Contains(strings.ToLower(err.Error()), "no longer exists") {
			// stream does not exist, may have been deleted accidentally
			a.logger.Info("Creating stream and consumer group", zap.String("group", groupName))
			//XGROUP CREATE creates the stream automatically, if it doesn't exist, when MKSTREAM subcommand is specified as last argument
			_, err := conn.Do("XGROUP", "CREATE", a.config.Stream, groupName, "$", "MKSTREAM")
			if err != nil {
				a.logger.Error("Cannot create stream and consumer group", zap.Error(err))
				return err
			}
		} else {
			return err
		}

	} else {

		if _, ok := groups[groupName]; ok {
			a.logger.Info("Reusing consumer group", zap.String("group", groupName))
		} else {
			a.logger.Info("Creating consumer group", zap.String("group", groupName))
			_, err := conn.Do("XGROUP", "CREATE", a.config.Stream, groupName, "$")
			if err != nil {
				a.logger.Error("Cannot create consumer group", zap.Error(err))
				return err
			}
		}

	}
	conn.Close()

	a.logger.Info("Number of consumers from config:", zap.Int("NumConsumers", a.config.NumConsumers))

	for i := 0; i < a.config.NumConsumers; i++ {
		go func(j int) {
			conn, _ := pool.Dial()

			consumerName := fmt.Sprintf("%s-%d", groupName, j)
			xreadID := "0" //Initial ID to read pending messages
			a.logger.Info("Listening for messages", zap.String("consumerName", consumerName))

			for {
				//XREAD below reads all the pending messages when xreadID=="0" and new messages when xreadID==">"
				reply, err := conn.Do("XREADGROUP", "GROUP", groupName, consumerName, "COUNT", 1, "BLOCK", 0, "STREAMS", a.config.Stream, xreadID)

				if err != nil {
					if strings.Contains(strings.ToLower(err.Error()), "no such key") || strings.Contains(strings.ToLower(err.Error()), "no longer exists") {
						a.logger.Info("Re-creating stream and consumer group", zap.String("group", groupName))
						_, err := conn.Do("XGROUP", "CREATE", a.config.Stream, groupName, "$", "MKSTREAM")
						if err != nil {
							a.logger.Error("Cannot re-create stream and group", zap.Error(err))
							time.Sleep(1 * time.Second)
							continue
						}
					} else {
						a.logger.Error("Cannot read from stream", zap.Error(err))
						time.Sleep(1 * time.Second)
						continue
					}
				}

				event, err := a.toEvent(reply)
				if err != nil {
					if strings.Contains(strings.ToLower(err.Error()), "number of items not equal to one (got 0)") { // no more pending messages!
						xreadID = ">" //ID to read new messages in next iteration
						continue
					} else {
						a.logger.Error("Cannot convert reply", zap.Error(err))
						continue
					}
				}
				xreadID = "0" //ID to read pending messages in next iteration

				if result := a.client.Send(ctx, *event); !cloudevents.IsACK(result) {
					//  Event is lost.
					a.logger.Error("Failed to send cloudevent", zap.Any("result", result))
				}
				//a.logger.Info("Consumer sent message as cloudevent", zap.String("consumerName", consumerName))

				_, err = conn.Do("XACK", a.config.Stream, groupName, event.ID())
				if err != nil {
					a.logger.Error("Cannot ack message", zap.Error(err))
				}
				a.logger.Info("Consumer acknowledged the message", zap.String("consumerName", consumerName))
			}
		}(i)
	}

	<-ctx.Done()
	return nil
}

func newPool(address string) *redis.Pool {
	return &redis.Pool{
		// Maximum number of idle connections in the pool.
		MaxIdle: 80,
		// max number of connections
		MaxActive: 12000,
		// Dial is an application supplied function for creating and
		// configuring a connection.
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", address)
			if err != nil {
				panic(err.Error())
			}
			return c, err
		},
	}
}

func (a *Adapter) toEvent(reply interface{}) (*cloudevents.Event, error) {
	values, err := redis.Values(reply, nil)
	if err != nil {
		return nil, errors.New("expected a reply of type array")
	}

	// Assert only one stream
	if len(values) != 1 {
		return nil, fmt.Errorf("number of values not equal to one (got %d)", len(values))
	}

	elems, err := scan.ScanXReadReply(values, nil)
	if err != nil {
		return nil, err
	}

	// Assert only one item
	if len(elems[0].Items) != 1 {
		return nil, fmt.Errorf("number of items not equal to one (got %d)", len(elems[0].Items))
	}

	item := elems[0].Items[0]

	event := cloudevents.NewEvent()
	event.SetType(RedisStreamSourceEventType)
	event.SetSource(a.source)
	event.SetData(cloudevents.ApplicationJSON, item.FieldValues)
	event.SetID(item.ID)

	return &event, nil
}
<commit_msg>Avoiding double xreads<commit_after>/*
Copyright 2020 The Knative Authors

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 adapter

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	scan "knative.dev/eventing-redis/source/pkg/redis"

	cloudevents "github.com/cloudevents/sdk-go/v2"
	"github.com/gomodule/redigo/redis"
	"go.uber.org/zap"
	"knative.dev/eventing/pkg/adapter/v2"
	"knative.dev/pkg/logging"
)

const (
	// RedisStreamSourceEventType is the default RedisStreamSource CloudEvent type.
	RedisStreamSourceEventType = "dev.knative.sources.redisstream"
)

func NewEnvConfig() adapter.EnvConfigAccessor {
	return &Config{}
}

type Adapter struct {
	config *Config
	logger *zap.Logger
	client cloudevents.Client
	source string
}

func NewAdapter(ctx context.Context, processed adapter.EnvConfigAccessor, ceClient cloudevents.Client) adapter.Adapter {
	config := processed.(*Config)
	return &Adapter{
		config: config,
		logger: logging.FromContext(ctx).Desugar().With(zap.String("stream", config.Stream)),
		client: ceClient,
		source: fmt.Sprintf("%s/%s", config.Address, config.Stream),
	}
}

func (a *Adapter) Start(ctx context.Context) error {
	pool := newPool(a.config.Address)

	conn, err := pool.Dial()
	if err != nil {
		return err
	}

	// Build consumer group name from stateful set pod name of adapter
	groupName := a.config.PodName

	a.logger.Info("Retrieving group info", zap.String("group", groupName))
	groups, err := scan.ScanXInfoGroupReply(conn.Do("XINFO", "GROUPS", a.config.Stream))

	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "no such key") || strings.Contains(strings.ToLower(err.Error()), "no longer exists") {
			// stream does not exist, may have been deleted accidentally
			a.logger.Info("Creating stream and consumer group", zap.String("group", groupName))
			//XGROUP CREATE creates the stream automatically, if it doesn't exist, when MKSTREAM subcommand is specified as last argument
			_, err := conn.Do("XGROUP", "CREATE", a.config.Stream, groupName, "$", "MKSTREAM")
			if err != nil {
				a.logger.Error("Cannot create stream and consumer group", zap.Error(err))
				return err
			}
		} else {
			return err
		}

	} else {

		if _, ok := groups[groupName]; ok {
			a.logger.Info("Reusing consumer group", zap.String("group", groupName))
		} else {
			a.logger.Info("Creating consumer group", zap.String("group", groupName))
			_, err := conn.Do("XGROUP", "CREATE", a.config.Stream, groupName, "$")
			if err != nil {
				a.logger.Error("Cannot create consumer group", zap.Error(err))
				return err
			}
		}

	}
	conn.Close()

	a.logger.Info("Number of consumers from config:", zap.Int("NumConsumers", a.config.NumConsumers))

	for i := 0; i < a.config.NumConsumers; i++ {
		go func(j int) {
			conn, _ := pool.Dial()

			consumerName := fmt.Sprintf("%s-%d", groupName, j)
			xreadID := "0" //Initial ID to read pending messages
			a.logger.Info("Listening for messages", zap.String("consumerName", consumerName))

			for {
				//XREAD below reads all the pending messages when xreadID=="0" and new messages when xreadID==">"
				reply, err := conn.Do("XREADGROUP", "GROUP", groupName, consumerName, "COUNT", 1, "BLOCK", 0, "STREAMS", a.config.Stream, xreadID)

				if err != nil {
					if strings.Contains(strings.ToLower(err.Error()), "no such key") || strings.Contains(strings.ToLower(err.Error()), "no longer exists") {
						a.logger.Info("Re-creating stream and consumer group", zap.String("group", groupName))
						_, err := conn.Do("XGROUP", "CREATE", a.config.Stream, groupName, "$", "MKSTREAM")
						if err != nil {
							a.logger.Error("Cannot re-create stream and group", zap.Error(err))
							time.Sleep(1 * time.Second)
							continue
						}
					} else {
						a.logger.Error("Cannot read from stream", zap.Error(err))
						time.Sleep(1 * time.Second)
						continue
					}
				}

				event, err := a.toEvent(reply)
				if err != nil {
					if strings.Contains(strings.ToLower(err.Error()), "number of items not equal to one (got 0)") { // no more pending messages!
						xreadID = ">" //ID to read new messages in next iteration
						continue
					} else {
						a.logger.Error("Cannot convert reply", zap.Error(err))
						time.Sleep(1 * time.Second)
						continue
					}
				}

				if result := a.client.Send(ctx, *event); !cloudevents.IsACK(result) {
					//  Event is lost.
					a.logger.Error("Failed to send cloudevent", zap.Any("result", result))
				}
				//a.logger.Info("Consumer sent message as cloudevent", zap.String("consumerName", consumerName))

				_, err = conn.Do("XACK", a.config.Stream, groupName, event.ID())
				if err != nil {
					a.logger.Error("Cannot ack message", zap.Error(err))
					xreadID = "0" //ID to read pending message in next iteration
					time.Sleep(1 * time.Second)
					continue
				}
				a.logger.Info("Consumer acknowledged the message", zap.String("consumerName", consumerName))
			}
		}(i)
	}

	<-ctx.Done()
	return nil
}

func newPool(address string) *redis.Pool {
	return &redis.Pool{
		// Maximum number of idle connections in the pool.
		MaxIdle: 80,
		// max number of connections
		MaxActive: 12000,
		// Dial is an application supplied function for creating and
		// configuring a connection.
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", address)
			if err != nil {
				panic(err.Error())
			}
			return c, err
		},
	}
}

func (a *Adapter) toEvent(reply interface{}) (*cloudevents.Event, error) {
	values, err := redis.Values(reply, nil)
	if err != nil {
		return nil, errors.New("expected a reply of type array")
	}

	// Assert only one stream
	if len(values) != 1 {
		return nil, fmt.Errorf("number of values not equal to one (got %d)", len(values))
	}

	elems, err := scan.ScanXReadReply(values, nil)
	if err != nil {
		return nil, err
	}

	// Assert only one item
	if len(elems[0].Items) != 1 {
		return nil, fmt.Errorf("number of items not equal to one (got %d)", len(elems[0].Items))
	}

	item := elems[0].Items[0]

	event := cloudevents.NewEvent()
	event.SetType(RedisStreamSourceEventType)
	event.SetSource(a.source)
	event.SetData(cloudevents.ApplicationJSON, item.FieldValues)
	event.SetID(item.ID)

	return &event, nil
}
<|endoftext|>


<commit_before>package libkbfs

import (
	"testing"

	"github.com/golang/mock/gomock"
	"golang.org/x/net/context"
)

func crTestInit(t *testing.T) (mockCtrl *gomock.Controller, config *ConfigMock,
	cr *ConflictResolver) {
	ctr := NewSafeTestReporter(t)
	mockCtrl = gomock.NewController(ctr)
	config = NewConfigMock(mockCtrl, ctr)
	id, _, _ := NewFolder(t, 1, 1, false, false)
	fbo := NewFolderBranchOps(config, FolderBranch{id, MasterBranch}, standard)
	cr = NewConflictResolver(config, fbo)
	return mockCtrl, config, cr
}

func crTestShutdown(mockCtrl *gomock.Controller, config *ConfigMock,
	cr *ConflictResolver) {
	config.ctr.CheckForFailures()
	cr.fbo.Shutdown()
	mockCtrl.Finish()
}

func TestCRInput(t *testing.T) {
	mockCtrl, config, cr := crTestInit(t)
	defer crTestShutdown(mockCtrl, config, cr)
	ctx := context.Background()

	// First try a completely unknown revision
	cr.Resolve(MetadataRevisionUninitialized, MetadataRevisionUninitialized)
	// This should return without doing anything (i.e., without
	// calling any mock methods)
	cr.Wait(ctx)

	// Next, try resolving a few items
	branchPoint := MetadataRevision(2)
	unmergedHead := MetadataRevision(5)
	mergedHead := MetadataRevision(15)

	cr.fbo.head = &RootMetadata{
		Revision: unmergedHead,
		Flags:    MetadataFlagUnmerged,
	}
	// serve all the MDs from the cache
	for i := unmergedHead; i >= branchPoint+1; i-- {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Unmerged).Return(
			&RootMetadata{
				Revision: i,
				Flags:    MetadataFlagUnmerged,
			}, nil)
	}
	for i := MetadataRevisionInitial; i <= branchPoint; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Unmerged).Return(
			nil, NoSuchMDError{cr.fbo.id(), branchPoint, Unmerged})
	}
	config.mockMdops.EXPECT().GetUnmergedRange(gomock.Any(), cr.fbo.id(),
		MetadataRevisionInitial, branchPoint).Return(nil, nil)

	for i := branchPoint + 1; i <= mergedHead; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
			&RootMetadata{Revision: i}, nil)
	}
	for i := mergedHead + 1; i <= branchPoint+2*maxMDsAtATime; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
			nil, NoSuchMDError{cr.fbo.id(), mergedHead + 1, Merged})
	}
	config.mockMdops.EXPECT().GetRange(gomock.Any(), cr.fbo.id(),
		mergedHead+1, gomock.Any()).Return(nil, nil)

	// First try a completely unknown revision
	cr.Resolve(unmergedHead, MetadataRevisionUninitialized)
	cr.Wait(ctx)
	// Make sure sure the input is up-to-date
	if cr.currInput.merged != mergedHead {
		t.Fatalf("Unexpected merged input: %d\n", cr.currInput.merged)
	}

	// Now make sure we ignore future inputs with lesser MDs
	cr.Resolve(MetadataRevisionUninitialized, mergedHead-1)
	// This should return without doing anything (i.e., without
	// calling any mock methods)
	cr.Wait(ctx)
}
<commit_msg>cr: add a test for when there's a hole in the cached MD range<commit_after>package libkbfs

import (
	"testing"

	"github.com/golang/mock/gomock"
	"golang.org/x/net/context"
)

func crTestInit(t *testing.T) (mockCtrl *gomock.Controller, config *ConfigMock,
	cr *ConflictResolver) {
	ctr := NewSafeTestReporter(t)
	mockCtrl = gomock.NewController(ctr)
	config = NewConfigMock(mockCtrl, ctr)
	id, _, _ := NewFolder(t, 1, 1, false, false)
	fbo := NewFolderBranchOps(config, FolderBranch{id, MasterBranch}, standard)
	cr = NewConflictResolver(config, fbo)
	return mockCtrl, config, cr
}

func crTestShutdown(mockCtrl *gomock.Controller, config *ConfigMock,
	cr *ConflictResolver) {
	config.ctr.CheckForFailures()
	cr.fbo.Shutdown()
	mockCtrl.Finish()
}

func TestCRInput(t *testing.T) {
	mockCtrl, config, cr := crTestInit(t)
	defer crTestShutdown(mockCtrl, config, cr)
	ctx := context.Background()

	// First try a completely unknown revision
	cr.Resolve(MetadataRevisionUninitialized, MetadataRevisionUninitialized)
	// This should return without doing anything (i.e., without
	// calling any mock methods)
	cr.Wait(ctx)

	// Next, try resolving a few items
	branchPoint := MetadataRevision(2)
	unmergedHead := MetadataRevision(5)
	mergedHead := MetadataRevision(15)

	cr.fbo.head = &RootMetadata{
		Revision: unmergedHead,
		Flags:    MetadataFlagUnmerged,
	}
	// serve all the MDs from the cache
	for i := unmergedHead; i >= branchPoint+1; i-- {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Unmerged).Return(
			&RootMetadata{
				Revision: i,
				Flags:    MetadataFlagUnmerged,
			}, nil)
	}
	for i := MetadataRevisionInitial; i <= branchPoint; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Unmerged).Return(
			nil, NoSuchMDError{cr.fbo.id(), branchPoint, Unmerged})
	}
	config.mockMdops.EXPECT().GetUnmergedRange(gomock.Any(), cr.fbo.id(),
		MetadataRevisionInitial, branchPoint).Return(nil, nil)

	for i := branchPoint + 1; i <= mergedHead; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
			&RootMetadata{Revision: i}, nil)
	}
	for i := mergedHead + 1; i <= branchPoint+2*maxMDsAtATime; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
			nil, NoSuchMDError{cr.fbo.id(), i, Merged})
	}
	config.mockMdops.EXPECT().GetRange(gomock.Any(), cr.fbo.id(),
		mergedHead+1, gomock.Any()).Return(nil, nil)

	// First try a completely unknown revision
	cr.Resolve(unmergedHead, MetadataRevisionUninitialized)
	cr.Wait(ctx)
	// Make sure sure the input is up-to-date
	if cr.currInput.merged != mergedHead {
		t.Fatalf("Unexpected merged input: %d\n", cr.currInput.merged)
	}

	// Now make sure we ignore future inputs with lesser MDs
	cr.Resolve(MetadataRevisionUninitialized, mergedHead-1)
	// This should return without doing anything (i.e., without
	// calling any mock methods)
	cr.Wait(ctx)
}

func TestCRInputFracturedRange(t *testing.T) {
	mockCtrl, config, cr := crTestInit(t)
	defer crTestShutdown(mockCtrl, config, cr)
	ctx := context.Background()

	// Next, try resolving a few items
	branchPoint := MetadataRevision(2)
	unmergedHead := MetadataRevision(5)
	mergedHead := MetadataRevision(15)

	cr.fbo.head = &RootMetadata{
		Revision: unmergedHead,
		Flags:    MetadataFlagUnmerged,
	}
	// serve all the MDs from the cache
	for i := unmergedHead; i >= branchPoint+1; i-- {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Unmerged).Return(
			&RootMetadata{
				Revision: i,
				Flags:    MetadataFlagUnmerged,
			}, nil)
	}
	for i := MetadataRevisionInitial; i <= branchPoint; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Unmerged).Return(
			nil, NoSuchMDError{cr.fbo.id(), branchPoint, Unmerged})
	}
	config.mockMdops.EXPECT().GetUnmergedRange(gomock.Any(), cr.fbo.id(),
		MetadataRevisionInitial, branchPoint).Return(nil, nil)

	skipCacheRevision := MetadataRevision(10)
	for i := branchPoint + 1; i <= mergedHead; i++ {
		// Pretend that revision 10 isn't in the cache, and needs to
		// be fetched from the server.
		if i != skipCacheRevision {
			config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
				&RootMetadata{Revision: i}, nil)
		} else {
			config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
				nil, NoSuchMDError{cr.fbo.id(), i, Merged})
		}
	}
	config.mockMdops.EXPECT().GetRange(gomock.Any(), cr.fbo.id(),
		skipCacheRevision, skipCacheRevision).Return(
		[]*RootMetadata{&RootMetadata{Revision: skipCacheRevision}}, nil)
	for i := mergedHead + 1; i <= branchPoint+2*maxMDsAtATime; i++ {
		config.mockMdcache.EXPECT().Get(cr.fbo.id(), i, Merged).Return(
			nil, NoSuchMDError{cr.fbo.id(), i, Merged})
	}
	config.mockMdops.EXPECT().GetRange(gomock.Any(), cr.fbo.id(),
		mergedHead+1, gomock.Any()).Return(nil, nil)

	// First try a completely unknown revision
	cr.Resolve(unmergedHead, MetadataRevisionUninitialized)
	cr.Wait(ctx)
	// Make sure sure the input is up-to-date
	if cr.currInput.merged != mergedHead {
		t.Fatalf("Unexpected merged input: %d\n", cr.currInput.merged)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Google 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 imageproxy provides an image proxy server.  For typical use of
// creating and using a Proxy, see cmd/imageproxy/main.go.
package imageproxy // import "willnorris.com/go/imageproxy"

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/gregjones/httpcache"
	tphttp "willnorris.com/go/imageproxy/third_party/http"
)

// Proxy serves image requests.
type Proxy struct {
	Client *http.Client // client used to fetch remote URLs
	Cache  Cache        // cache used to cache responses

	// Whitelist specifies a list of remote hosts that images can be
	// proxied from.  An empty list means all hosts are allowed.
	Whitelist []string

	// Referrers, when given, requires that requests to the image
	// proxy come from a referring host. An empty list means all
	// hosts are allowed.
	Referrers []string

	// DefaultBaseURL is the URL that relative remote URLs are resolved in
	// reference to.  If nil, all remote URLs specified in requests must be
	// absolute.
	DefaultBaseURL *url.URL

	// SignatureKey is the HMAC key used to verify signed requests.
	SignatureKey []byte

	// Allow images to scale beyond their original dimensions.
	ScaleUp bool

	// Timeout specifies a time limit for requests served by this Proxy.
	// If a call runs for longer than its time limit, a 504 Gateway Timeout
	// response is returned.  A Timeout of zero means no timeout.
	Timeout time.Duration
}

// NewProxy constructs a new proxy.  The provided http RoundTripper will be
// used to fetch remote URLs.  If nil is provided, http.DefaultTransport will
// be used.
func NewProxy(transport http.RoundTripper, cache Cache) *Proxy {
	if transport == nil {
		transport = http.DefaultTransport
	}
	if cache == nil {
		cache = NopCache
	}

	proxy := Proxy{
		Cache: cache,
	}

	client := new(http.Client)
	client.Transport = &httpcache.Transport{
		Transport:           &TransformingTransport{transport, client},
		Cache:               cache,
		MarkCachedResponses: true,
	}

	proxy.Client = client

	return &proxy
}

// ServeHTTP handles incoming requests.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/favicon.ico" {
		return // ignore favicon requests
	}

	if r.URL.Path == "/health-check" {
		fmt.Fprint(w, "OK")
		return
	}

	var h http.Handler = http.HandlerFunc(p.serveImage)
	if p.Timeout > 0 {
		h = tphttp.TimeoutHandler(h, p.Timeout, "Gateway timeout waiting for remote resource.")
	}
	h.ServeHTTP(w, r)
}

// serveImage handles incoming requests for proxied images.
func (p *Proxy) serveImage(w http.ResponseWriter, r *http.Request) {
	req, err := NewRequest(r, p.DefaultBaseURL)
	if err != nil {
		msg := fmt.Sprintf("invalid request URL: %v", err)
		glog.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	// assign static settings from proxy to req.Options
	req.Options.ScaleUp = p.ScaleUp

	if err := p.allowed(req); err != nil {
		glog.Error(err)
		http.Error(w, err.Error(), http.StatusForbidden)
		return
	}

	resp, err := p.Client.Get(req.String())
	if err != nil {
		msg := fmt.Sprintf("error fetching remote image: %v", err)
		glog.Error(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	cached := resp.Header.Get(httpcache.XFromCache)
	glog.Infof("request: %v (served from cache: %v)", *req, cached == "1")

	copyHeader(w.Header(), resp.Header, "Cache-Control", "Last-Modified", "Expires", "Etag", "Link")

	if should304(r, resp) {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	copyHeader(w.Header(), resp.Header, "Content-Length", "Content-Type")
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

// copyHeader copies header values from src to dst, adding to any existing
// values with the same header name.  If keys is not empty, only those header
// keys will be copied.
func copyHeader(dst, src http.Header, keys ...string) {
	if len(keys) == 0 {
		for k, _ := range src {
			keys = append(keys, k)
		}
	}
	for _, key := range keys {
		k := http.CanonicalHeaderKey(key)
		for _, v := range src[k] {
			dst.Add(k, v)
		}
	}
}

// allowed determines whether the specified request contains an allowed
// referrer, host, and signature.  It returns an error if the request is not
// allowed.
func (p *Proxy) allowed(r *Request) error {
	if len(p.Referrers) > 0 && !validReferrer(p.Referrers, r.Original) {
		return fmt.Errorf("request does not contain an allowed referrer: %v", r)
	}

	if len(p.Whitelist) == 0 && len(p.SignatureKey) == 0 {
		return nil // no whitelist or signature key, all requests accepted
	}

	if len(p.Whitelist) > 0 && validHost(p.Whitelist, r.URL) {
		return nil
	}

	if len(p.SignatureKey) > 0 && validSignature(p.SignatureKey, r) {
		return nil
	}

	return fmt.Errorf("request does not contain an allowed host or valid signature: %v", r)
}

// validHost returns whether the host in u matches one of hosts.
func validHost(hosts []string, u *url.URL) bool {
	for _, host := range hosts {
		if u.Host == host {
			return true
		}
		if strings.HasPrefix(host, "*.") && strings.HasSuffix(u.Host, host[2:]) {
			return true
		}
	}

	return false
}

// returns whether the referrer from the request is in the host list.
func validReferrer(hosts []string, r *http.Request) bool {
	u, err := url.Parse(r.Header.Get("Referer"))
	if err != nil { // malformed or blank header, just deny
		return false
	}

	return validHost(hosts, u)
}

// validSignature returns whether the request signature is valid.
func validSignature(key []byte, r *Request) bool {
	sig := r.Options.Signature
	if m := len(sig) % 4; m != 0 { // add padding if missing
		sig += strings.Repeat("=", 4-m)
	}

	got, err := base64.URLEncoding.DecodeString(sig)
	if err != nil {
		glog.Errorf("error base64 decoding signature %q", r.Options.Signature)
		return false
	}

	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(r.URL.String()))
	want := mac.Sum(nil)

	return hmac.Equal(got, want)
}

// should304 returns whether we should send a 304 Not Modified in response to
// req, based on the response resp.  This is determined using the last modified
// time and the entity tag of resp.
func should304(req *http.Request, resp *http.Response) bool {
	// TODO(willnorris): if-none-match header can be a comma separated list
	// of multiple tags to be matched, or the special value "*" which
	// matches all etags
	etag := resp.Header.Get("Etag")
	if etag != "" && etag == req.Header.Get("If-None-Match") {
		return true
	}

	lastModified, err := time.Parse(time.RFC1123, resp.Header.Get("Last-Modified"))
	if err != nil {
		return false
	}
	ifModSince, err := time.Parse(time.RFC1123, req.Header.Get("If-Modified-Since"))
	if err != nil {
		return false
	}
	if lastModified.Before(ifModSince) {
		return true
	}

	return false
}

// TransformingTransport is an implementation of http.RoundTripper that
// optionally transforms images using the options specified in the request URL
// fragment.
type TransformingTransport struct {
	// Transport is the underlying http.RoundTripper used to satisfy
	// non-transform requests (those that do not include a URL fragment).
	Transport http.RoundTripper

	// CachingClient is used to fetch images to be resized.  This client is
	// used rather than Transport directly in order to ensure that
	// responses are properly cached.
	CachingClient *http.Client
}

// RoundTrip implements the http.RoundTripper interface.
func (t *TransformingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.URL.Fragment == "" {
		// normal requests pass through
		glog.Infof("fetching remote URL: %v", req.URL)
		return t.Transport.RoundTrip(req)
	}

	u := *req.URL
	u.Fragment = ""
	resp, err := t.CachingClient.Get(u.String())
	if err != nil {
		return nil, err
	}

	if should304(req, resp) {
		// bare 304 response, full response will be used from cache
		return &http.Response{StatusCode: http.StatusNotModified}, nil
	}

	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	opt := ParseOptions(req.URL.Fragment)

	img, err := Transform(b, opt)
	if err != nil {
		glog.Errorf("error transforming image: %v", err)
		img = b
	}

	// replay response with transformed image and updated content length
	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "%s %s\n", resp.Proto, resp.Status)
	resp.Header.WriteSubset(buf, map[string]bool{
		"Content-Length": true,
		// exclude Content-Type header if the format may have changed during transformation
		"Content-Type": opt.Format != "" || resp.Header.Get("Content-Type") == "image/webp",
	})
	fmt.Fprintf(buf, "Content-Length: %d\n\n", len(img))
	buf.Write(img)

	return http.ReadResponse(bufio.NewReader(buf), req)
}
<commit_msg>Enable CORS for 3rd party applications<commit_after>// Copyright 2013 Google 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 imageproxy provides an image proxy server.  For typical use of
// creating and using a Proxy, see cmd/imageproxy/main.go.
package imageproxy // import "willnorris.com/go/imageproxy"

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/gregjones/httpcache"
	tphttp "willnorris.com/go/imageproxy/third_party/http"
)

// Proxy serves image requests.
type Proxy struct {
	Client *http.Client // client used to fetch remote URLs
	Cache  Cache        // cache used to cache responses

	// Whitelist specifies a list of remote hosts that images can be
	// proxied from.  An empty list means all hosts are allowed.
	Whitelist []string

	// Referrers, when given, requires that requests to the image
	// proxy come from a referring host. An empty list means all
	// hosts are allowed.
	Referrers []string

	// DefaultBaseURL is the URL that relative remote URLs are resolved in
	// reference to.  If nil, all remote URLs specified in requests must be
	// absolute.
	DefaultBaseURL *url.URL

	// SignatureKey is the HMAC key used to verify signed requests.
	SignatureKey []byte

	// Allow images to scale beyond their original dimensions.
	ScaleUp bool

	// Timeout specifies a time limit for requests served by this Proxy.
	// If a call runs for longer than its time limit, a 504 Gateway Timeout
	// response is returned.  A Timeout of zero means no timeout.
	Timeout time.Duration
}

// NewProxy constructs a new proxy.  The provided http RoundTripper will be
// used to fetch remote URLs.  If nil is provided, http.DefaultTransport will
// be used.
func NewProxy(transport http.RoundTripper, cache Cache) *Proxy {
	if transport == nil {
		transport = http.DefaultTransport
	}
	if cache == nil {
		cache = NopCache
	}

	proxy := Proxy{
		Cache: cache,
	}

	client := new(http.Client)
	client.Transport = &httpcache.Transport{
		Transport:           &TransformingTransport{transport, client},
		Cache:               cache,
		MarkCachedResponses: true,
	}

	proxy.Client = client

	return &proxy
}

// ServeHTTP handles incoming requests.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/favicon.ico" {
		return // ignore favicon requests
	}

	if r.URL.Path == "/health-check" {
		fmt.Fprint(w, "OK")
		return
	}

	var h http.Handler = http.HandlerFunc(p.serveImage)
	if p.Timeout > 0 {
		h = tphttp.TimeoutHandler(h, p.Timeout, "Gateway timeout waiting for remote resource.")
	}
	h.ServeHTTP(w, r)
}

// serveImage handles incoming requests for proxied images.
func (p *Proxy) serveImage(w http.ResponseWriter, r *http.Request) {
	req, err := NewRequest(r, p.DefaultBaseURL)
	if err != nil {
		msg := fmt.Sprintf("invalid request URL: %v", err)
		glog.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	// assign static settings from proxy to req.Options
	req.Options.ScaleUp = p.ScaleUp

	if err := p.allowed(req); err != nil {
		glog.Error(err)
		http.Error(w, err.Error(), http.StatusForbidden)
		return
	}

	resp, err := p.Client.Get(req.String())
	if err != nil {
		msg := fmt.Sprintf("error fetching remote image: %v", err)
		glog.Error(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	cached := resp.Header.Get(httpcache.XFromCache)
	glog.Infof("request: %v (served from cache: %v)", *req, cached == "1")

	copyHeader(w.Header(), resp.Header, "Cache-Control", "Last-Modified", "Expires", "Etag", "Link")

	if should304(r, resp) {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	copyHeader(w.Header(), resp.Header, "Content-Length", "Content-Type")

	//Enable CORS for 3rd party applications
	w.Header().Set("Access-Control-Allow-Origin", "*")

	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

// copyHeader copies header values from src to dst, adding to any existing
// values with the same header name.  If keys is not empty, only those header
// keys will be copied.
func copyHeader(dst, src http.Header, keys ...string) {
	if len(keys) == 0 {
		for k, _ := range src {
			keys = append(keys, k)
		}
	}
	for _, key := range keys {
		k := http.CanonicalHeaderKey(key)
		for _, v := range src[k] {
			dst.Add(k, v)
		}
	}
}

// allowed determines whether the specified request contains an allowed
// referrer, host, and signature.  It returns an error if the request is not
// allowed.
func (p *Proxy) allowed(r *Request) error {
	if len(p.Referrers) > 0 && !validReferrer(p.Referrers, r.Original) {
		return fmt.Errorf("request does not contain an allowed referrer: %v", r)
	}

	if len(p.Whitelist) == 0 && len(p.SignatureKey) == 0 {
		return nil // no whitelist or signature key, all requests accepted
	}

	if len(p.Whitelist) > 0 && validHost(p.Whitelist, r.URL) {
		return nil
	}

	if len(p.SignatureKey) > 0 && validSignature(p.SignatureKey, r) {
		return nil
	}

	return fmt.Errorf("request does not contain an allowed host or valid signature: %v", r)
}

// validHost returns whether the host in u matches one of hosts.
func validHost(hosts []string, u *url.URL) bool {
	for _, host := range hosts {
		if u.Host == host {
			return true
		}
		if strings.HasPrefix(host, "*.") && strings.HasSuffix(u.Host, host[2:]) {
			return true
		}
	}

	return false
}

// returns whether the referrer from the request is in the host list.
func validReferrer(hosts []string, r *http.Request) bool {
	u, err := url.Parse(r.Header.Get("Referer"))
	if err != nil { // malformed or blank header, just deny
		return false
	}

	return validHost(hosts, u)
}

// validSignature returns whether the request signature is valid.
func validSignature(key []byte, r *Request) bool {
	sig := r.Options.Signature
	if m := len(sig) % 4; m != 0 { // add padding if missing
		sig += strings.Repeat("=", 4-m)
	}

	got, err := base64.URLEncoding.DecodeString(sig)
	if err != nil {
		glog.Errorf("error base64 decoding signature %q", r.Options.Signature)
		return false
	}

	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(r.URL.String()))
	want := mac.Sum(nil)

	return hmac.Equal(got, want)
}

// should304 returns whether we should send a 304 Not Modified in response to
// req, based on the response resp.  This is determined using the last modified
// time and the entity tag of resp.
func should304(req *http.Request, resp *http.Response) bool {
	// TODO(willnorris): if-none-match header can be a comma separated list
	// of multiple tags to be matched, or the special value "*" which
	// matches all etags
	etag := resp.Header.Get("Etag")
	if etag != "" && etag == req.Header.Get("If-None-Match") {
		return true
	}

	lastModified, err := time.Parse(time.RFC1123, resp.Header.Get("Last-Modified"))
	if err != nil {
		return false
	}
	ifModSince, err := time.Parse(time.RFC1123, req.Header.Get("If-Modified-Since"))
	if err != nil {
		return false
	}
	if lastModified.Before(ifModSince) {
		return true
	}

	return false
}

// TransformingTransport is an implementation of http.RoundTripper that
// optionally transforms images using the options specified in the request URL
// fragment.
type TransformingTransport struct {
	// Transport is the underlying http.RoundTripper used to satisfy
	// non-transform requests (those that do not include a URL fragment).
	Transport http.RoundTripper

	// CachingClient is used to fetch images to be resized.  This client is
	// used rather than Transport directly in order to ensure that
	// responses are properly cached.
	CachingClient *http.Client
}

// RoundTrip implements the http.RoundTripper interface.
func (t *TransformingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.URL.Fragment == "" {
		// normal requests pass through
		glog.Infof("fetching remote URL: %v", req.URL)
		return t.Transport.RoundTrip(req)
	}

	u := *req.URL
	u.Fragment = ""
	resp, err := t.CachingClient.Get(u.String())
	if err != nil {
		return nil, err
	}

	if should304(req, resp) {
		// bare 304 response, full response will be used from cache
		return &http.Response{StatusCode: http.StatusNotModified}, nil
	}

	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	opt := ParseOptions(req.URL.Fragment)

	img, err := Transform(b, opt)
	if err != nil {
		glog.Errorf("error transforming image: %v", err)
		img = b
	}

	// replay response with transformed image and updated content length
	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "%s %s\n", resp.Proto, resp.Status)
	resp.Header.WriteSubset(buf, map[string]bool{
		"Content-Length": true,
		// exclude Content-Type header if the format may have changed during transformation
		"Content-Type": opt.Format != "" || resp.Header.Get("Content-Type") == "image/webp",
	})
	fmt.Fprintf(buf, "Content-Length: %d\n\n", len(img))
	buf.Write(img)

	return http.ReadResponse(bufio.NewReader(buf), req)
}
<|endoftext|>