<commit_before>package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dustin/go-id3"
	"github.com/dustin/goexif/exif"
)

var uploadWg = sync.WaitGroup{}

var uploadFlags = flag.NewFlagSet("upload", flag.ExitOnError)
var uploadVerbose = uploadFlags.Bool("v", false, "Verbose")
var uploadDelete = uploadFlags.Bool("delete", false,
	"Delete locally missing items")
var uploadMeta = uploadFlags.Bool("meta", false,
	"Store meta info in userData for items")
var uploadWorkers = uploadFlags.Int("workers", 4, "Number of upload workers")
var uploadRevs = uploadFlags.Int("revs", 0,
	"Number of old revisions to keep (-1 == all)")
var uploadIgnore = uploadFlags.String("ignore", "",
	"Path to ignore file")
var uploadRevsSet = false

var quotingReplacer = strings.NewReplacer("%", "%25",
	"?", "%3f",
	" ", "%20",
	"#", "%23")

type uploadOpType uint8

const (
	uploadFileOp = uploadOpType(iota)
	removeFileOp
	removeRecurseOp
)

type uploadReq struct {
	src        string
	dest       string
	op         uploadOpType
	remoteHash string
}

func recognizeTypeByName(n, def string) string {
	byname := mime.TypeByExtension(n)
	switch {
	case byname != "":
		return byname
	case strings.HasSuffix(n, ".js"):
		return "application/javascript"
	case strings.HasSuffix(n, ".css"):
		return "text/css"
	}
	return def
}

func processMP3Meta(src, dest string) (interface{}, error) {
	f, err := os.Open(src)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	ifile := id3.Read(f)
	return ifile, nil
}

func processEXIFMeta(src, dest string) (interface{}, error) {
	f, err := os.Open(src)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	return exif.Decode(f)
}

func processMeta(src, dest string) error {
	var data interface{}
	var err error

	switch filepath.Ext(strings.ToLower(src)) {
	case ".mp3":
		data, err = processMP3Meta(src, dest)
	case ".jpg", ".jpeg":
		data, err = processEXIFMeta(src, dest)
	default:
		log.Printf("No meta info for %#v",
			filepath.Ext(strings.ToLower(src)))
	}

	if err != nil || data == nil {
		return err
	}

	b, err := json.Marshal(data)
	if err != nil {
		return err
	}

	udest, err := url.Parse(dest)
	if err != nil {
		return err
	}
	udest.Path = "/.cbfs/meta" + udest.Path

	preq, err := http.NewRequest("PUT", udest.String(), bytes.NewReader(b))
	if err != nil {
		return err
	}

	preq.Header.Set("Content-Type", "application/json")

	if *uploadVerbose {
		log.Printf("Uploading meta info to %v", udest)
	}

	resp, err := http.DefaultClient.Do(preq)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 201 {
		return fmt.Errorf("HTTP Error:  %v", resp.Status)
	}

	return nil
}

func uploadFile(src, dest, localHash string) error {
	if *uploadVerbose {
		log.Printf("Uploading %v -> %v (%v)", src, dest, localHash)
	}

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

	someBytes := make([]byte, 512)
	n, err := f.Read(someBytes)
	if err != nil && err != io.EOF {
		return err
	}
	someBytes = someBytes[:n]

	length, err := f.Seek(0, 2)
	if err != nil {
		return err
	}

	_, err = f.Seek(0, 0)
	if err != nil {
		return err
	}

	preq, err := http.NewRequest("PUT", dest, f)
	if err != nil {
		return err
	}
	if uploadRevsSet {
		preq.Header.Set("X-CBFS-KeepRevs", strconv.Itoa(*uploadRevs))
	}

	ctype := http.DetectContentType(someBytes)
	if strings.HasPrefix(ctype, "text/plain") ||
		strings.HasPrefix(ctype, "application/octet-stream") {
		ctype = recognizeTypeByName(src, ctype)
	}

	preq.Header.Set("Content-Length", strconv.FormatInt(length, 10))
	preq.Header.Set("Content-Type", ctype)
	preq.Header.Set("X-CBFS-Hash", localHash)

	resp, err := http.DefaultClient.Do(preq)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 201 {
		r, _ := ioutil.ReadAll(resp.Body)
		return fmt.Errorf("HTTP Error:  %v: %s", resp.Status, r)
	}

	if *uploadMeta {
		err = processMeta(src, dest)
		if err != nil {
			log.Printf("Error processing meta info: %v", err)
		}
	}

	return nil
}

// This is very similar to rm's version, but uses different channel
// signaling.
func uploadRmDashR(baseUrl string, ch chan uploadReq) ([]string, error) {
	r := quotingReplacer
	for strings.HasSuffix(baseUrl, "/") {
		baseUrl = baseUrl[:len(baseUrl)-1]
	}

	listing, err := listOrEmpty(baseUrl)
	for err != nil {
		return []string{}, err
	}
	for fn := range listing.Files {
		err = rmFile(baseUrl + "/" + r.Replace(fn))
		if err != nil {
			return []string{}, err
		}
	}
	children := make([]string, 0, len(listing.Dirs))
	for dn := range listing.Dirs {
		children = append(children, baseUrl+"/"+r.Replace(dn))
	}
	return children, nil
}

func localHash(fn string) string {
	f, err := os.Open(fn)
	if err != nil {
		return "unknown"
	}
	defer f.Close()

	h := sha1.New()
	_, err = io.Copy(h, f)
	if err != nil {
		return "unknown"
	}

	return hex.EncodeToString(h.Sum([]byte{}))
}

func uploadWorker(ch chan uploadReq) {
	defer uploadWg.Done()
	for req := range ch {
		retries := 0
		done := false
		for !done {
			var err error
			switch req.op {
			case uploadFileOp:
				lh := localHash(req.src)
				if req.remoteHash == "" {
					err = uploadFile(req.src, req.dest, lh)
				} else {
					if lh != req.remoteHash {
						if *uploadVerbose {
							log.Printf("%v has changed, reupping",
								req.src)
						}
						err = uploadFile(req.src, req.dest, lh)
					}
				}
			case removeFileOp:
				if *uploadVerbose {
					log.Printf("Removing file %v", req.dest)
				}
				err = rmFile(req.dest)
			case removeRecurseOp:
				todo := []string{req.dest}
				for err == nil && len(todo) > 0 {
					todo, err = uploadRmDashR(req.dest, ch)
				}
			default:
				log.Fatalf("Unhandled case")
			}
			if err != nil {
				if retries < 3 {
					retries++
					log.Printf("Error uploading file: %v... retrying",
						err)
					time.Sleep(time.Duration(retries) * time.Second)
				} else {
					log.Printf("Error uploading file %v: %v",
						req.src, err)
					done = true
				}
			} else {
				done = true
			}
		}
	}
}

func syncPath(path, dest string, info os.FileInfo, ch chan<- uploadReq) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	children, err := f.Readdir(0)
	if err != nil {
		return err
	}

	dest = quotingReplacer.Replace(dest)

	retries := 3
	serverListing, err := listOrEmpty(dest)
	for err != nil && retries > 0 {
		serverListing, err = listOrEmpty(dest)
		time.Sleep(time.Second)
		retries--
	}
	if err != nil {
		return err
	}

	localNames := map[string]os.FileInfo{}
	for _, c := range children {
		switch c.Mode() & os.ModeType {
		case os.ModeCharDevice, os.ModeDevice,
			os.ModeNamedPipe, os.ModeSocket, os.ModeSymlink:
			if *uploadVerbose {
				log.Printf("Ignoring special file: %v - %v",
					filepath.Join(path, c.Name()), c.Mode())
			}
		default:
			fullPath := filepath.Join(path, c.Name())
			if !isIgnored(fullPath) {
				localNames[c.Name()] = c
			} else {
				if *uploadVerbose {
					log.Printf("Ignoring %v", fullPath)
				}
			}
		}
	}

	remoteNames := map[string]bool{}
	for n := range serverListing.Files {
		if n != "" {
			remoteNames[n] = true
		}
	}
	for n := range serverListing.Dirs {
		if n != "" {
			remoteNames[n] = true
		}
	}

	// Keeping it short
	r := quotingReplacer

	missingUpstream := []string{}
	for n, fi := range localNames {
		if !(fi.IsDir() || remoteNames[n]) {
			missingUpstream = append(missingUpstream, n)
		} else if !fi.IsDir() {
			if ri, ok := serverListing.Files[n]; ok {
				ch <- uploadReq{filepath.Join(path, n),
					dest + "/" + r.Replace(n), uploadFileOp, ri.OID}
			}
		}
	}

	toRm := []string{}
	for n := range remoteNames {
		if _, ok := localNames[n]; !ok {
			toRm = append(toRm, n)
		}
	}

	if len(missingUpstream) > 0 {
		for _, m := range missingUpstream {
			ch <- uploadReq{filepath.Join(path, m),
				dest + "/" + r.Replace(m), uploadFileOp, ""}
		}
	}

	if *uploadDelete && len(toRm) > 0 {
		for _, m := range toRm {
			ch <- uploadReq{"", dest + "/" + r.Replace(m), removeFileOp, ""}
			ch <- uploadReq{"", dest + "/" + r.Replace(m), removeRecurseOp, ""}
		}
	}

	return nil
}

func syncUp(src, u string, ch chan<- uploadReq) {
	for strings.HasSuffix(u, "/") {
		u = u[:len(u)-1]
	}
	for strings.HasSuffix(src, "/") {
		src = src[:len(src)-1]
	}

	err := filepath.Walk(src,
		func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				if isIgnored(path) {
					log.Printf("Skipping directory %v",
						path)
					return filepath.SkipDir
				}
				shortPath := path[len(src):]
				err = syncPath(path, u+shortPath, info, ch)
			}
			return err
		})
	if err != nil {
		log.Fatalf("Traversal error: %v", err)
	}
}

func uploadCommand(u string, args []string) {
	uploadFlags.Parse(args)

	uploadFlags.Visit(func(f *flag.Flag) {
		if f.Name == "revs" {
			uploadRevsSet = true
		}
	})

	if uploadFlags.NArg() < 2 {
		log.Fatalf("src and dest required")
	}

	if *uploadIgnore != "" {
		err := loadIgnorePatternsFromFile(*uploadIgnore)
		if err != nil {
			log.Fatalf("Error loading ignores: %v", err)
		}
	}

	du := relativeUrl(u, uploadFlags.Arg(1))

	fi, err := os.Stat(uploadFlags.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	if fi.IsDir() {
		ch := make(chan uploadReq, 1000)

		for i := 0; i < *uploadWorkers; i++ {
			uploadWg.Add(1)
			go uploadWorker(ch)
		}

		start := time.Now()
		syncUp(uploadFlags.Arg(0), du, ch)

		close(ch)
		log.Printf("Finished traversal in %v", time.Since(start))
		uploadWg.Wait()
		log.Printf("Finished sync in %v", time.Since(start))
	} else {
		err = uploadFile(uploadFlags.Arg(0), du,
			localHash(uploadFlags.Arg(0)))
		if err != nil {
			log.Fatalf("Error uploading file: %v", err)
		}
	}
}
<commit_msg>Log the upload actions a bit more accurately.<commit_after>package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dustin/go-id3"
	"github.com/dustin/goexif/exif"
)

var uploadWg = sync.WaitGroup{}

var uploadFlags = flag.NewFlagSet("upload", flag.ExitOnError)
var uploadVerbose = uploadFlags.Bool("v", false, "Verbose")
var uploadDelete = uploadFlags.Bool("delete", false,
	"Delete locally missing items")
var uploadMeta = uploadFlags.Bool("meta", false,
	"Store meta info in userData for items")
var uploadWorkers = uploadFlags.Int("workers", 4, "Number of upload workers")
var uploadRevs = uploadFlags.Int("revs", 0,
	"Number of old revisions to keep (-1 == all)")
var uploadIgnore = uploadFlags.String("ignore", "",
	"Path to ignore file")
var uploadRevsSet = false

var quotingReplacer = strings.NewReplacer("%", "%25",
	"?", "%3f",
	" ", "%20",
	"#", "%23")

type uploadOpType uint8

const (
	uploadFileOp = uploadOpType(iota)
	removeFileOp
	removeRecurseOp
)

func (u uploadOpType) String() string {
	switch u {
	case uploadFileOp:
		return "upload file"
	case removeFileOp:
		return "remove file"
	case removeRecurseOp:
		return "remove (recursive) file"
	}
	panic("unhandled op type")
}

type uploadReq struct {
	src        string
	dest       string
	op         uploadOpType
	remoteHash string
}

func recognizeTypeByName(n, def string) string {
	byname := mime.TypeByExtension(n)
	switch {
	case byname != "":
		return byname
	case strings.HasSuffix(n, ".js"):
		return "application/javascript"
	case strings.HasSuffix(n, ".css"):
		return "text/css"
	}
	return def
}

func processMP3Meta(src, dest string) (interface{}, error) {
	f, err := os.Open(src)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	ifile := id3.Read(f)
	return ifile, nil
}

func processEXIFMeta(src, dest string) (interface{}, error) {
	f, err := os.Open(src)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	return exif.Decode(f)
}

func processMeta(src, dest string) error {
	var data interface{}
	var err error

	switch filepath.Ext(strings.ToLower(src)) {
	case ".mp3":
		data, err = processMP3Meta(src, dest)
	case ".jpg", ".jpeg":
		data, err = processEXIFMeta(src, dest)
	default:
		log.Printf("No meta info for %#v",
			filepath.Ext(strings.ToLower(src)))
	}

	if err != nil || data == nil {
		return err
	}

	b, err := json.Marshal(data)
	if err != nil {
		return err
	}

	udest, err := url.Parse(dest)
	if err != nil {
		return err
	}
	udest.Path = "/.cbfs/meta" + udest.Path

	preq, err := http.NewRequest("PUT", udest.String(), bytes.NewReader(b))
	if err != nil {
		return err
	}

	preq.Header.Set("Content-Type", "application/json")

	if *uploadVerbose {
		log.Printf("Uploading meta info to %v", udest)
	}

	resp, err := http.DefaultClient.Do(preq)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 201 {
		return fmt.Errorf("HTTP Error:  %v", resp.Status)
	}

	return nil
}

func uploadFile(src, dest, localHash string) error {
	if *uploadVerbose {
		log.Printf("Uploading %v -> %v (%v)", src, dest, localHash)
	}

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

	someBytes := make([]byte, 512)
	n, err := f.Read(someBytes)
	if err != nil && err != io.EOF {
		return err
	}
	someBytes = someBytes[:n]

	length, err := f.Seek(0, 2)
	if err != nil {
		return err
	}

	_, err = f.Seek(0, 0)
	if err != nil {
		return err
	}

	preq, err := http.NewRequest("PUT", dest, f)
	if err != nil {
		return err
	}
	if uploadRevsSet {
		preq.Header.Set("X-CBFS-KeepRevs", strconv.Itoa(*uploadRevs))
	}

	ctype := http.DetectContentType(someBytes)
	if strings.HasPrefix(ctype, "text/plain") ||
		strings.HasPrefix(ctype, "application/octet-stream") {
		ctype = recognizeTypeByName(src, ctype)
	}

	preq.Header.Set("Content-Length", strconv.FormatInt(length, 10))
	preq.Header.Set("Content-Type", ctype)
	preq.Header.Set("X-CBFS-Hash", localHash)

	resp, err := http.DefaultClient.Do(preq)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 201 {
		r, _ := ioutil.ReadAll(resp.Body)
		return fmt.Errorf("HTTP Error:  %v: %s", resp.Status, r)
	}

	if *uploadMeta {
		err = processMeta(src, dest)
		if err != nil {
			log.Printf("Error processing meta info: %v", err)
		}
	}

	return nil
}

// This is very similar to rm's version, but uses different channel
// signaling.
func uploadRmDashR(baseUrl string, ch chan uploadReq) ([]string, error) {
	r := quotingReplacer
	for strings.HasSuffix(baseUrl, "/") {
		baseUrl = baseUrl[:len(baseUrl)-1]
	}

	listing, err := listOrEmpty(baseUrl)
	for err != nil {
		return []string{}, err
	}
	for fn := range listing.Files {
		err = rmFile(baseUrl + "/" + r.Replace(fn))
		if err != nil {
			return []string{}, err
		}
	}
	children := make([]string, 0, len(listing.Dirs))
	for dn := range listing.Dirs {
		children = append(children, baseUrl+"/"+r.Replace(dn))
	}
	return children, nil
}

func localHash(fn string) string {
	f, err := os.Open(fn)
	if err != nil {
		return "unknown"
	}
	defer f.Close()

	h := sha1.New()
	_, err = io.Copy(h, f)
	if err != nil {
		return "unknown"
	}

	return hex.EncodeToString(h.Sum([]byte{}))
}

func uploadWorker(ch chan uploadReq) {
	defer uploadWg.Done()
	for req := range ch {
		retries := 0
		done := false
		for !done {
			var err error
			switch req.op {
			case uploadFileOp:
				lh := localHash(req.src)
				if req.remoteHash == "" {
					err = uploadFile(req.src, req.dest, lh)
				} else {
					if lh != req.remoteHash {
						if *uploadVerbose {
							log.Printf("%v has changed, reupping",
								req.src)
						}
						err = uploadFile(req.src, req.dest, lh)
					}
				}
			case removeFileOp:
				if *uploadVerbose {
					log.Printf("Removing file %v", req.dest)
				}
				err = rmFile(req.dest)
			case removeRecurseOp:
				todo := []string{req.dest}
				for err == nil && len(todo) > 0 {
					todo, err = uploadRmDashR(req.dest, ch)
				}
			default:
				log.Fatalf("Unhandled case")
			}
			if err != nil {
				if retries < 3 {
					retries++
					log.Printf("Error in %v: %v... retrying",
						req.op, err)
					time.Sleep(time.Duration(retries) * time.Second)
				} else {
					log.Printf("Error in %v %v: %v",
						req.op, req.src, err)
					done = true
				}
			} else {
				done = true
			}
		}
	}
}

func syncPath(path, dest string, info os.FileInfo, ch chan<- uploadReq) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	children, err := f.Readdir(0)
	if err != nil {
		return err
	}

	dest = quotingReplacer.Replace(dest)

	retries := 3
	serverListing, err := listOrEmpty(dest)
	for err != nil && retries > 0 {
		serverListing, err = listOrEmpty(dest)
		time.Sleep(time.Second)
		retries--
	}
	if err != nil {
		return err
	}

	localNames := map[string]os.FileInfo{}
	for _, c := range children {
		switch c.Mode() & os.ModeType {
		case os.ModeCharDevice, os.ModeDevice,
			os.ModeNamedPipe, os.ModeSocket, os.ModeSymlink:
			if *uploadVerbose {
				log.Printf("Ignoring special file: %v - %v",
					filepath.Join(path, c.Name()), c.Mode())
			}
		default:
			fullPath := filepath.Join(path, c.Name())
			if !isIgnored(fullPath) {
				localNames[c.Name()] = c
			} else {
				if *uploadVerbose {
					log.Printf("Ignoring %v", fullPath)
				}
			}
		}
	}

	remoteNames := map[string]bool{}
	for n := range serverListing.Files {
		if n != "" {
			remoteNames[n] = true
		}
	}
	for n := range serverListing.Dirs {
		if n != "" {
			remoteNames[n] = true
		}
	}

	// Keeping it short
	r := quotingReplacer

	missingUpstream := []string{}
	for n, fi := range localNames {
		if !(fi.IsDir() || remoteNames[n]) {
			missingUpstream = append(missingUpstream, n)
		} else if !fi.IsDir() {
			if ri, ok := serverListing.Files[n]; ok {
				ch <- uploadReq{filepath.Join(path, n),
					dest + "/" + r.Replace(n), uploadFileOp, ri.OID}
			}
		}
	}

	toRm := []string{}
	for n := range remoteNames {
		if _, ok := localNames[n]; !ok {
			toRm = append(toRm, n)
		}
	}

	if len(missingUpstream) > 0 {
		for _, m := range missingUpstream {
			ch <- uploadReq{filepath.Join(path, m),
				dest + "/" + r.Replace(m), uploadFileOp, ""}
		}
	}

	if *uploadDelete && len(toRm) > 0 {
		for _, m := range toRm {
			ch <- uploadReq{"", dest + "/" + r.Replace(m), removeFileOp, ""}
			ch <- uploadReq{"", dest + "/" + r.Replace(m), removeRecurseOp, ""}
		}
	}

	return nil
}

func syncUp(src, u string, ch chan<- uploadReq) {
	for strings.HasSuffix(u, "/") {
		u = u[:len(u)-1]
	}
	for strings.HasSuffix(src, "/") {
		src = src[:len(src)-1]
	}

	err := filepath.Walk(src,
		func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				if isIgnored(path) {
					log.Printf("Skipping directory %v",
						path)
					return filepath.SkipDir
				}
				shortPath := path[len(src):]
				err = syncPath(path, u+shortPath, info, ch)
			}
			return err
		})
	if err != nil {
		log.Fatalf("Traversal error: %v", err)
	}
}

func uploadCommand(u string, args []string) {
	uploadFlags.Parse(args)

	uploadFlags.Visit(func(f *flag.Flag) {
		if f.Name == "revs" {
			uploadRevsSet = true
		}
	})

	if uploadFlags.NArg() < 2 {
		log.Fatalf("src and dest required")
	}

	if *uploadIgnore != "" {
		err := loadIgnorePatternsFromFile(*uploadIgnore)
		if err != nil {
			log.Fatalf("Error loading ignores: %v", err)
		}
	}

	du := relativeUrl(u, uploadFlags.Arg(1))

	fi, err := os.Stat(uploadFlags.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	if fi.IsDir() {
		ch := make(chan uploadReq, 1000)

		for i := 0; i < *uploadWorkers; i++ {
			uploadWg.Add(1)
			go uploadWorker(ch)
		}

		start := time.Now()
		syncUp(uploadFlags.Arg(0), du, ch)

		close(ch)
		log.Printf("Finished traversal in %v", time.Since(start))
		uploadWg.Wait()
		log.Printf("Finished sync in %v", time.Since(start))
	} else {
		err = uploadFile(uploadFlags.Arg(0), du,
			localHash(uploadFlags.Arg(0)))
		if err != nil {
			log.Fatalf("Error uploading file: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package avroio

import (
	"context"
	"encoding/json"
	"reflect"
	"strings"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/io/filesystem"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	"github.com/linkedin/goavro"
)

// Avroio package based directly on textio with some
// elements borrowed from the bigquery implementation.
// Supports reading and unmarshalling Avro files and
// returns either RAW JSON string or Unmarshalled type
// --

func init() {
	// beam.RegisterType(reflect.TypeOf((*writeFileFn)(nil)).Elem())
	// beam.RegisterFunction(readFn)
	beam.RegisterFunction(expandFn)
	beam.RegisterType(reflect.TypeOf((*avroReadFn)(nil)).Elem())
}

// Read reads a set of files and returns the lines as a PCollection<elem>
// based on avro schema. Support a type via < reflect.TypeOf(YourType{}) >  with
// JSON tags defined or if you wish to return the raw JSON string, use < reflect.TypeOf("") >
func Read(s beam.Scope, glob string, t reflect.Type) beam.PCollection {
	s = s.Scope("avroio.Read")
	filesystem.ValidateScheme(glob)
	return read(s, t, beam.Create(s, glob))
}

func read(s beam.Scope, t reflect.Type, col beam.PCollection) beam.PCollection {
	files := beam.ParDo(s, expandFn, col)
	return beam.ParDo(s,
		&avroReadFn{Type: beam.EncodedType{T: t}},
		files,
		beam.TypeDefinition{Var: beam.XType, T: t},
	)
}

func expandFn(ctx context.Context, glob string, emit func(string)) error {
	if strings.TrimSpace(glob) == "" {
		return nil // ignore empty string elements here
	}

	fs, err := filesystem.New(ctx, glob)
	if err != nil {
		return err
	}
	defer fs.Close()

	files, err := fs.List(ctx, glob)
	if err != nil {
		return err
	}
	for _, filename := range files {
		emit(filename)
	}
	return nil
}

type avroReadFn struct {
	// Avro schema type
	Type beam.EncodedType
}

func (f *avroReadFn) ProcessElement(ctx context.Context, filename string, emit func(beam.X)) (err error) {
	log.Infof(ctx, "Reading from %v", filename)

	fs, err := filesystem.New(ctx, filename)
	if err != nil {
		return
	}
	defer fs.Close()

	fd, err := fs.OpenRead(ctx, filename)
	if err != nil {
		return
	}
	defer fd.Close()

	ar, err := goavro.NewOCFReader(fd)
	if err != nil {
		log.Errorf(ctx, "error reading avro: %v", err)
		return
	}

	val := reflect.New(f.Type.T).Interface()
	for ar.Scan() {
		var i interface{}
		i, err = ar.Read()
		if err != nil {
			log.Errorf(ctx, "error reading avro row: %v", err)
			continue
		}

		// marshal interface to bytes
		var b []byte
		b, err = json.Marshal(i)
		if err != nil {
			log.Errorf(ctx, "error unmarshalling avro data: %v", err)
			return
		}

		switch reflect.New(f.Type.T).Interface().(type) {
		case *string:
			emit(string(b))
		default:
			if err = json.Unmarshal(b, val); err != nil {
				log.Errorf(ctx, "error unmashalling avro to type: %v", err)
				return
			}
			emit(reflect.ValueOf(val).Elem().Interface())
		}
	}

	return ar.Err()
}

// TODO: Implement Avroio Write Functionanilty
<commit_msg>added avro write support<commit_after>package avroio

import (
	"context"
	"encoding/json"
	"reflect"
	"strings"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/io/filesystem"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	"github.com/linkedin/goavro"
)

// Avroio package based directly on textio with some
// elements borrowed from the bigquery implementation.
// Supports reading/writing and unmarshalling Avro files.
// --

func init() {
	// beam.RegisterType(reflect.TypeOf((*writeFileFn)(nil)).Elem())
	// beam.RegisterFunction(readFn)
	beam.RegisterFunction(expandFn)
	beam.RegisterType(reflect.TypeOf((*avroReadFn)(nil)).Elem())
}

// Read reads a set of files and returns the lines as a PCollection<elem>
// based on avro schema. Support a type via < reflect.TypeOf(YourType{}) >  with
// JSON tags defined or if you wish to return the raw JSON string, use < reflect.TypeOf("") >
func Read(s beam.Scope, glob string, t reflect.Type) beam.PCollection {
	s = s.Scope("avroio.Read")
	filesystem.ValidateScheme(glob)
	return read(s, t, beam.Create(s, glob))
}

func read(s beam.Scope, t reflect.Type, col beam.PCollection) beam.PCollection {
	files := beam.ParDo(s, expandFn, col)
	return beam.ParDo(s,
		&avroReadFn{Type: beam.EncodedType{T: t}},
		files,
		beam.TypeDefinition{Var: beam.XType, T: t},
	)
}

func expandFn(ctx context.Context, glob string, emit func(string)) error {
	if strings.TrimSpace(glob) == "" {
		return nil // ignore empty string elements here
	}

	fs, err := filesystem.New(ctx, glob)
	if err != nil {
		return err
	}
	defer fs.Close()

	files, err := fs.List(ctx, glob)
	if err != nil {
		return err
	}
	for _, filename := range files {
		emit(filename)
	}
	return nil
}

type avroReadFn struct {
	// Avro schema type
	Type beam.EncodedType
}

func (f *avroReadFn) ProcessElement(ctx context.Context, filename string, emit func(beam.X)) (err error) {
	log.Infof(ctx, "Reading from %v", filename)

	fs, err := filesystem.New(ctx, filename)
	if err != nil {
		return
	}
	defer fs.Close()

	fd, err := fs.OpenRead(ctx, filename)
	if err != nil {
		return
	}
	defer fd.Close()

	ar, err := goavro.NewOCFReader(fd)
	if err != nil {
		log.Errorf(ctx, "error reading avro: %v", err)
		return
	}

	val := reflect.New(f.Type.T).Interface()
	for ar.Scan() {
		var i interface{}
		i, err = ar.Read()
		if err != nil {
			log.Errorf(ctx, "error reading avro row: %v", err)
			continue
		}

		// marshal interface to bytes
		var b []byte
		b, err = json.Marshal(i)
		if err != nil {
			log.Errorf(ctx, "error unmarshalling avro data: %v", err)
			return
		}

		switch reflect.New(f.Type.T).Interface().(type) {
		case *string:
			emit(string(b))
		default:
			if err = json.Unmarshal(b, val); err != nil {
				log.Errorf(ctx, "error unmashalling avro to type: %v", err)
				return
			}
			emit(reflect.ValueOf(val).Elem().Interface())
		}
	}

	return ar.Err()
}

// Write writes a PCollection<string> to an AVRO file.
// Write expects a JSON string with a matching AVRO schema.
// the process will fail if the schema does not match the JSON
// provided
func Write(s beam.Scope, filename, schema string, col beam.PCollection) {
	s = s.Scope("textio.Write")
	filesystem.ValidateScheme(filename)
	pre := beam.AddFixedKey(s, col)
	post := beam.GroupByKey(s, pre)
	beam.ParDo0(s, &writeAvroFn{Schema: schema, Filename: filename}, post)
}

type writeAvroFn struct {
	Schema   string `json:"schema"`
	Filename string `json:"filename"`
}

func (w *writeAvroFn) ProcessElement(ctx context.Context, _ int, lines func(*string) bool) (err error) {
	fs, err := filesystem.New(ctx, w.Filename)
	if err != nil {
		return
	}
	defer fs.Close()

	fd, err := fs.OpenWrite(ctx, w.Filename)
	if err != nil {
		return
	}

	defer fd.Close()

	codec, err := goavro.NewCodec(w.Schema)
	if err != nil {
		log.Errorf(ctx, "error creating avro codec: %v", err)
		return
	}

	ocfw, err := goavro.NewOCFWriter(goavro.OCFConfig{
		Codec:           codec,
		CompressionName: goavro.CompressionSnappyLabel,
		Schema:          w.Schema,
		W:               fd,
	})

	if err != nil {
		log.Errorf(ctx, "error creating avro writer: %v", err)
		return
	}

	var j string
	for lines(&j) {
		native, _, err := codec.NativeFromTextual([]byte(j))
		if err != nil {
			log.Errorf(ctx, "error reading native avro: %v", err)
			return err
		}

		if err := ocfw.Append([]interface{}{native}); err != nil {
			log.Errorf(ctx, "error writing avro: %v", err)
			return err
		}
	}

	return
}
<|endoftext|>


<commit_before>package proxy

import (
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"

	log "github.com/sirupsen/logrus"

	"bytes"

	"github.com/gorilla/websocket"
)

var (
	// DefaultUpgrader specifies the parameters for upgrading an HTTP
	// connection to a WebSocket connection.
	DefaultUpgrader = &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	// DefaultDialer is a dialer with all fields set to the default zero values.
	DefaultDialer = websocket.DefaultDialer
)

// WebsocketProxy is an HTTP Handler that takes an incoming WebSocket
// connection and proxies it to another server.
type WebsocketProxy struct {
	// Director, if non-nil, is a function that may copy additional request
	// headers from the incoming WebSocket connection into the output headers
	// which will be forwarded to another server.
	Director func(incoming *http.Request, out http.Header)

	// Backend returns the backend URL which the proxy uses to reverse proxy
	// the incoming WebSocket connection. Request is the initial incoming and
	// unmodified request.
	Backend func(*http.Request) *url.URL

	// Upgrader specifies the parameters for upgrading a incoming HTTP
	// connection to a WebSocket connection. If nil, DefaultUpgrader is used.
	Upgrader *websocket.Upgrader

	//  Dialer contains options for connecting to the backend WebSocket server.
	//  If nil, DefaultDialer is used.
	Dialer *websocket.Dialer
}

// ProxyHandler returns a new http.Handler interface that reverse proxies the
// request to the given target.
func WebsocketProxyHandler(target *url.URL) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		NewWebsocketProxy(target).ServeHTTP(rw, req)
	})
}

// NewProxy returns a new Websocket reverse proxy that rewrites the
// URL's to the scheme, host and base path provider in target.
func NewWebsocketProxy(target *url.URL) *WebsocketProxy {
	backend := func(r *http.Request) *url.URL {
		// Shallow copy
		u := *target
		u.Fragment = r.URL.Fragment
		u.Path = r.URL.Path
		u.RawQuery = r.URL.RawQuery
		return &u
	}
	return &WebsocketProxy{Backend: backend}
}

// ServeHTTP implements the http.Handler that proxies WebSocket connections.
func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if w.Backend == nil {
		log.Error("WebsocketProxy: backend function is not defined")
		http.Error(rw, "internal server error (code: 1)", http.StatusInternalServerError)
		return
	}

	backendURL := w.Backend(req)
	if backendURL == nil {
		log.Error("WebsocketProxy: backend URL is nil")
		http.Error(rw, "internal server error (code: 2)", http.StatusInternalServerError)
		return
	}

	dialer := w.Dialer
	if w.Dialer == nil {
		dialer = DefaultDialer
	}

	// Pass headers from the incoming request to the dialer to forward them to
	// the final destinations.
	requestHeader := http.Header{}
	if origin := req.Header.Get("Origin"); origin != "" {
		requestHeader.Add("Origin", origin)
	}
	for _, prot := range req.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")] {
		requestHeader.Add("Sec-WebSocket-Protocol", prot)
	}
	for _, cookie := range req.Header[http.CanonicalHeaderKey("Cookie")] {
		requestHeader.Add("Cookie", cookie)
	}

	// Pass X-Forwarded-For headers too, code below is a part of
	// httputil.ReverseProxy. See http://en.wikipedia.org/wiki/X-Forwarded-For
	// for more information
	// TODO: use RFC7239 http://tools.ietf.org/html/rfc7239
	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := req.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		requestHeader.Set("X-Forwarded-For", clientIP)
	}

	// Set the originating protocol of the incoming HTTP request. The SSL might
	// be terminated on our site and because we doing proxy adding this would
	// be helpful for applications on the backend.
	requestHeader.Set("X-Forwarded-Proto", "http")
	if req.TLS != nil {
		requestHeader.Set("X-Forwarded-Proto", "https")
	}

	// Enable the director to copy any additional headers it desires for
	// forwarding to the remote server.
	if w.Director != nil {
		w.Director(req, requestHeader)
	}

	// Connect to the backend URL, also pass the headers we get from the requst
	// together with the Forwarded headers we prepared above.
	// TODO: support multiplexing on the same backend connection instead of
	// opening a new TCP connection time for each request. This should be
	// optional:
	// http://tools.ietf.org/html/draft-ietf-hybi-websocket-multiplexing-01
	connBackend, resp, err := dialer.Dial(backendURL.String(), requestHeader)
	if err != nil {
		buff := &bytes.Buffer{}
		buff.ReadFrom(resp.Body)
		log.Errorf("WebsocketProxy: couldn't dial to remote backend url %s: %s %s %s", backendURL, err, resp.Status, buff.String())
		return
	}
	defer connBackend.Close()

	upgrader := w.Upgrader
	if w.Upgrader == nil {
		upgrader = DefaultUpgrader
	}

	// Only pass those headers to the upgrader.
	upgradeHeader := http.Header{}
	if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {
		upgradeHeader.Set("Sec-Websocket-Protocol", hdr)
	}
	if hdr := resp.Header.Get("Set-Cookie"); hdr != "" {
		upgradeHeader.Set("Set-Cookie", hdr)
	}

	// Now upgrade the existing incoming request to a WebSocket connection.
	// Also pass the header that we gathered from the Dial handshake.
	connPub, err := upgrader.Upgrade(rw, req, upgradeHeader)
	if err != nil {
		log.Printf("websocketproxy: couldn't upgrade %s\n", err)
		return
	}
	defer connPub.Close()

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}

	// Start our proxy now, everything is ready...
	go cp(connBackend.UnderlyingConn(), connPub.UnderlyingConn())
	go cp(connPub.UnderlyingConn(), connBackend.UnderlyingConn())
	<-errc
}
<commit_msg>fix for possible nil pointer on failed ws connection<commit_after>package proxy

import (
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"

	log "github.com/sirupsen/logrus"

	"bytes"

	"github.com/gorilla/websocket"
)

var (
	// DefaultUpgrader specifies the parameters for upgrading an HTTP
	// connection to a WebSocket connection.
	DefaultUpgrader = &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	// DefaultDialer is a dialer with all fields set to the default zero values.
	DefaultDialer = websocket.DefaultDialer
)

// WebsocketProxy is an HTTP Handler that takes an incoming WebSocket
// connection and proxies it to another server.
type WebsocketProxy struct {
	// Director, if non-nil, is a function that may copy additional request
	// headers from the incoming WebSocket connection into the output headers
	// which will be forwarded to another server.
	Director func(incoming *http.Request, out http.Header)

	// Backend returns the backend URL which the proxy uses to reverse proxy
	// the incoming WebSocket connection. Request is the initial incoming and
	// unmodified request.
	Backend func(*http.Request) *url.URL

	// Upgrader specifies the parameters for upgrading a incoming HTTP
	// connection to a WebSocket connection. If nil, DefaultUpgrader is used.
	Upgrader *websocket.Upgrader

	//  Dialer contains options for connecting to the backend WebSocket server.
	//  If nil, DefaultDialer is used.
	Dialer *websocket.Dialer
}

// ProxyHandler returns a new http.Handler interface that reverse proxies the
// request to the given target.
func WebsocketProxyHandler(target *url.URL) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		NewWebsocketProxy(target).ServeHTTP(rw, req)
	})
}

// NewProxy returns a new Websocket reverse proxy that rewrites the
// URL's to the scheme, host and base path provider in target.
func NewWebsocketProxy(target *url.URL) *WebsocketProxy {
	backend := func(r *http.Request) *url.URL {
		// Shallow copy
		u := *target
		u.Fragment = r.URL.Fragment
		u.Path = r.URL.Path
		u.RawQuery = r.URL.RawQuery
		return &u
	}
	return &WebsocketProxy{Backend: backend}
}

// ServeHTTP implements the http.Handler that proxies WebSocket connections.
func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if w.Backend == nil {
		log.Error("WebsocketProxy: backend function is not defined")
		http.Error(rw, "internal server error (code: 1)", http.StatusInternalServerError)
		return
	}

	backendURL := w.Backend(req)
	if backendURL == nil {
		log.Error("WebsocketProxy: backend URL is nil")
		http.Error(rw, "internal server error (code: 2)", http.StatusInternalServerError)
		return
	}

	dialer := w.Dialer
	if w.Dialer == nil {
		dialer = DefaultDialer
	}

	// Pass headers from the incoming request to the dialer to forward them to
	// the final destinations.
	requestHeader := http.Header{}
	if origin := req.Header.Get("Origin"); origin != "" {
		requestHeader.Add("Origin", origin)
	}
	for _, prot := range req.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")] {
		requestHeader.Add("Sec-WebSocket-Protocol", prot)
	}
	for _, cookie := range req.Header[http.CanonicalHeaderKey("Cookie")] {
		requestHeader.Add("Cookie", cookie)
	}

	// Pass X-Forwarded-For headers too, code below is a part of
	// httputil.ReverseProxy. See http://en.wikipedia.org/wiki/X-Forwarded-For
	// for more information
	// TODO: use RFC7239 http://tools.ietf.org/html/rfc7239
	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := req.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		requestHeader.Set("X-Forwarded-For", clientIP)
	}

	// Set the originating protocol of the incoming HTTP request. The SSL might
	// be terminated on our site and because we doing proxy adding this would
	// be helpful for applications on the backend.
	requestHeader.Set("X-Forwarded-Proto", "http")
	if req.TLS != nil {
		requestHeader.Set("X-Forwarded-Proto", "https")
	}

	// Enable the director to copy any additional headers it desires for
	// forwarding to the remote server.
	if w.Director != nil {
		w.Director(req, requestHeader)
	}

	// Connect to the backend URL, also pass the headers we get from the requst
	// together with the Forwarded headers we prepared above.
	// TODO: support multiplexing on the same backend connection instead of
	// opening a new TCP connection time for each request. This should be
	// optional:
	// http://tools.ietf.org/html/draft-ietf-hybi-websocket-multiplexing-01
	connBackend, resp, err := dialer.Dial(backendURL.String(), requestHeader)
	if err != nil {
		status := "<none>"
		body := ""
		if resp != nil {
			status = resp.Status
			if resp.Body != nil {
				buff := &bytes.Buffer{}
				buff.ReadFrom(resp.Body)
				body = buff.String()
			}
		}
		log.Errorf("WebsocketProxy: couldn't dial to remote backend url %s: %s %s %s", backendURL, err, status, body)
		return
	}
	defer connBackend.Close()

	upgrader := w.Upgrader
	if w.Upgrader == nil {
		upgrader = DefaultUpgrader
	}

	// Only pass those headers to the upgrader.
	upgradeHeader := http.Header{}
	if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {
		upgradeHeader.Set("Sec-Websocket-Protocol", hdr)
	}
	if hdr := resp.Header.Get("Set-Cookie"); hdr != "" {
		upgradeHeader.Set("Set-Cookie", hdr)
	}

	// Now upgrade the existing incoming request to a WebSocket connection.
	// Also pass the header that we gathered from the Dial handshake.
	connPub, err := upgrader.Upgrade(rw, req, upgradeHeader)
	if err != nil {
		log.Printf("websocketproxy: couldn't upgrade %s\n", err)
		return
	}
	defer connPub.Close()

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}

	// Start our proxy now, everything is ready...
	go cp(connBackend.UnderlyingConn(), connPub.UnderlyingConn())
	go cp(connPub.UnderlyingConn(), connBackend.UnderlyingConn())
	<-errc
}
<|endoftext|>


<commit_before>package vfs

import (
	"errors"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
)

// ErrIteratorDone is returned by the Next() merthod of the iterator when
// the iterator is actually done.
var ErrIteratorDone = errors.New("No more element in the iterator")

// IteratorDefaultFetchSize is the default number of elements fetched from
// couchdb on each iteration.
const IteratorDefaultFetchSize = 100

// IteratorOptions contains the options of the iterator.
type IteratorOptions struct {
	StartKey string
	ByFetch  int
}

// Iterator is a struct allowing to iterate over the children of a directory.
// The iterator is not thread-safe.
type Iterator struct {
	ctx    Context
	sel    mango.Filter
	opt    *IteratorOptions
	list   []*DirOrFileDoc
	offset int
	index  int
	done   bool
}

// NewIterator return a new iterator.
func NewIterator(c Context, sel mango.Filter, opt *IteratorOptions) *Iterator {
	if opt == nil {
		opt = &IteratorOptions{ByFetch: IteratorDefaultFetchSize}
	}
	if opt.ByFetch == 0 {
		opt.ByFetch = IteratorDefaultFetchSize
	}
	return &Iterator{
		ctx: c,
		sel: sel,
		opt: opt,
	}
}

// Next should be called to get the next directory or file children of the
// parent directory. If the error is ErrIteratorDone
func (i *Iterator) Next() (*DirDoc, *FileDoc, error) {
	if i.done {
		return nil, nil, ErrIteratorDone
	}
	if i.index >= len(i.list) {
		if err := i.fetch(); err != nil {
			return nil, nil, err
		}
	}
	d, f := i.list[i.index].Refine()
	i.index++
	return d, f, nil
}

// fetch should be called when the index is out of the list boundary.
func (i *Iterator) fetch() error {
	l := len(i.list)
	if l > 0 && l < i.opt.ByFetch {
		i.done = true
		return ErrIteratorDone
	}

	i.offset += l
	i.index = 0
	i.list = i.list[:0]

	var skip int
	sel := i.sel
	if i.opt.StartKey != "" {
		// TODO: adapt this code when filtering and sorting are added to the
		// iterator
		sel = mango.And(sel, mango.AfterID(i.opt.StartKey))
	} else {
		skip = i.offset
	}

	req := &couchdb.FindRequest{
		Selector: sel,
		Limit:    i.opt.ByFetch,
		Skip:     skip,
	}
	err := couchdb.FindDocs(i.ctx, consts.Files, req, &i.list)
	if err != nil {
		return err
	}
	if len(i.list) == 0 {
		return ErrIteratorDone
	}
	return nil
}
<commit_msg>Fix skipping in iterator<commit_after>package vfs

import (
	"errors"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
)

// ErrIteratorDone is returned by the Next() merthod of the iterator when
// the iterator is actually done.
var ErrIteratorDone = errors.New("No more element in the iterator")

// IteratorDefaultFetchSize is the default number of elements fetched from
// couchdb on each iteration.
const IteratorDefaultFetchSize = 100

// IteratorOptions contains the options of the iterator.
type IteratorOptions struct {
	StartKey string
	ByFetch  int
}

// Iterator is a struct allowing to iterate over the children of a directory.
// The iterator is not thread-safe.
type Iterator struct {
	ctx    Context
	sel    mango.Filter
	opt    *IteratorOptions
	list   []*DirOrFileDoc
	offset int
	index  int
	done   bool
}

// NewIterator return a new iterator.
func NewIterator(c Context, sel mango.Filter, opt *IteratorOptions) *Iterator {
	if opt == nil {
		opt = &IteratorOptions{ByFetch: IteratorDefaultFetchSize}
	}
	if opt.ByFetch == 0 {
		opt.ByFetch = IteratorDefaultFetchSize
	}
	return &Iterator{
		ctx: c,
		sel: sel,
		opt: opt,
	}
}

// Next should be called to get the next directory or file children of the
// parent directory. If the error is ErrIteratorDone
func (i *Iterator) Next() (*DirDoc, *FileDoc, error) {
	if i.done {
		return nil, nil, ErrIteratorDone
	}
	if i.index >= len(i.list) {
		if err := i.fetch(); err != nil {
			return nil, nil, err
		}
	}
	d, f := i.list[i.index].Refine()
	i.index++
	return d, f, nil
}

// fetch should be called when the index is out of the list boundary.
func (i *Iterator) fetch() error {
	l := len(i.list)
	if l > 0 && l < i.opt.ByFetch {
		i.done = true
		return ErrIteratorDone
	}

	i.offset += l
	i.index = 0
	i.list = i.list[:0]

	sel := i.sel
	if i.opt.StartKey != "" {
		// TODO: adapt this code when filtering and sorting are added to the
		// iterator
		sel = mango.And(sel, mango.AfterID(i.opt.StartKey))
	}

	req := &couchdb.FindRequest{
		Selector: sel,
		Limit:    i.opt.ByFetch,
		Skip:     i.offset,
	}
	err := couchdb.FindDocs(i.ctx, consts.Files, req, &i.list)
	if err != nil {
		return err
	}
	if len(i.list) == 0 {
		return ErrIteratorDone
	}
	return nil
}
<|endoftext|>


<commit_before>package txtdirect

import (
	"net/http"
	"net/url"
	"path"
	"strings"
)

func parsePlaceholders(input string, r *http.Request) string {
	placeholders := PlaceholderRegex.FindAllStringSubmatch(input, -1)
	for _, placeholder := range placeholders {
		switch placeholder[0] {
		case "{uri}":
			input = strings.Replace(input, "{uri}", r.URL.RequestURI(), -1)
		case "{dir}":
			dir, _ := path.Split(r.URL.Path)
			input = strings.Replace(input, "{dir}", dir, -1)
		case "{file}":
			_, file := path.Split(r.URL.Path)
			input = strings.Replace(input, "{file}", file, -1)
		case "{fragment}":
			input = strings.Replace(input, "{fragment}", r.URL.Fragment, -1)
		case "{host}":
			input = strings.Replace(input, "{host}", r.URL.Host, -1)
		case "{hostonly}":
			input = strings.Replace(input, "{hostonly}", r.URL.Hostname(), -1)
		case "{method}":
			input = strings.Replace(input, "{method}", r.Method, -1)
		case "{path}":
			input = strings.Replace(input, "{path}", r.URL.Path, -1)
		case "{path_escaped}":
			input = strings.Replace(input, "{path_escaped}", url.QueryEscape(r.URL.Path), -1)
		case "{port}":
			input = strings.Replace(input, "{port}", r.URL.Port(), -1)
		case "{query}":
			input = strings.Replace(input, "{query}", r.URL.RawQuery, -1)
		case "{query_escaped}":
			input = strings.Replace(input, "{query_escaped}", url.QueryEscape(r.URL.RawQuery), -1)
		case "{uri_escaped}":
			input = strings.Replace(input, "{uri_escaped}", url.QueryEscape(r.URL.RequestURI()), -1)
		case "{user}":
			user, _, ok := r.BasicAuth()
			if !ok {
				input = strings.Replace(input, "{user}", "", -1)
			}
			input = strings.Replace(input, "{user}", user, -1)
		}
		if placeholder[0][1] == '>' {
			want := placeholder[0][2 : len(placeholder[0])-1]
			for key, values := range r.Header {
				// Header placeholders (case-insensitive)
				if strings.EqualFold(key, want) {
					input = strings.Replace(input, placeholder[0], strings.Join(values, ","), -1)
				}
			}
		}
		if placeholder[0][1] == '~' {
			name := placeholder[0][2 : len(placeholder[0])-1]
			if cookie, err := r.Cookie(name); err == nil {
				input = strings.Replace(input, placeholder[0], cookie.Value, -1)
			}
		}
		if placeholder[0][1] == '?' {
			query := r.URL.Query()
			name := placeholder[0][2 : len(placeholder[0])-1]
			input = strings.Replace(input, placeholder[0], query.Get(name), -1)
		}
	}
	return input
}
<commit_msg>Add parsing for {labelN} placeholder<commit_after>package txtdirect

import (
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
)

func parsePlaceholders(input string, r *http.Request) string {
	placeholders := PlaceholderRegex.FindAllStringSubmatch(input, -1)
	for _, placeholder := range placeholders {
		switch placeholder[0] {
		case "{uri}":
			input = strings.Replace(input, "{uri}", r.URL.RequestURI(), -1)
		case "{dir}":
			dir, _ := path.Split(r.URL.Path)
			input = strings.Replace(input, "{dir}", dir, -1)
		case "{file}":
			_, file := path.Split(r.URL.Path)
			input = strings.Replace(input, "{file}", file, -1)
		case "{fragment}":
			input = strings.Replace(input, "{fragment}", r.URL.Fragment, -1)
		case "{host}":
			input = strings.Replace(input, "{host}", r.URL.Host, -1)
		case "{hostonly}":
			input = strings.Replace(input, "{hostonly}", r.URL.Hostname(), -1)
		case "{label":
			nStr := placeholder[0][6 : len(placeholder[0])-1] // get the integer N in "{labelN}"
			n, err := strconv.Atoi(nStr)
			if err != nil || n < 1 {
				input = strings.Replace(input, placeholder[0], "", -1)
			}
			labels := strings.Split(r.Host, ".")
			if n > len(labels) {
				input = strings.Replace(input, placeholder[0], "", -1)
			}
			input = strings.Replace(input, placeholder[0], labels[n-1], -1)
		case "{method}":
			input = strings.Replace(input, "{method}", r.Method, -1)
		case "{path}":
			input = strings.Replace(input, "{path}", r.URL.Path, -1)
		case "{path_escaped}":
			input = strings.Replace(input, "{path_escaped}", url.QueryEscape(r.URL.Path), -1)
		case "{port}":
			input = strings.Replace(input, "{port}", r.URL.Port(), -1)
		case "{query}":
			input = strings.Replace(input, "{query}", r.URL.RawQuery, -1)
		case "{query_escaped}":
			input = strings.Replace(input, "{query_escaped}", url.QueryEscape(r.URL.RawQuery), -1)
		case "{uri_escaped}":
			input = strings.Replace(input, "{uri_escaped}", url.QueryEscape(r.URL.RequestURI()), -1)
		case "{user}":
			user, _, ok := r.BasicAuth()
			if !ok {
				input = strings.Replace(input, "{user}", "", -1)
			}
			input = strings.Replace(input, "{user}", user, -1)
		}
		if placeholder[0][1] == '>' {
			want := placeholder[0][2 : len(placeholder[0])-1]
			for key, values := range r.Header {
				// Header placeholders (case-insensitive)
				if strings.EqualFold(key, want) {
					input = strings.Replace(input, placeholder[0], strings.Join(values, ","), -1)
				}
			}
		}
		if placeholder[0][1] == '~' {
			name := placeholder[0][2 : len(placeholder[0])-1]
			if cookie, err := r.Cookie(name); err == nil {
				input = strings.Replace(input, placeholder[0], cookie.Value, -1)
			}
		}
		if placeholder[0][1] == '?' {
			query := r.URL.Query()
			name := placeholder[0][2 : len(placeholder[0])-1]
			input = strings.Replace(input, placeholder[0], query.Get(name), -1)
		}
	}
	return input
}
<|endoftext|>


<commit_before>package player

func computeScore(keywords1, keywords2 []string) float64 {
	var total float64 = 0

	for _, k1 := range keywords1 {
		for _, k2 := range keywords2 {
			total = total + inverseLevenshteinDistance(k1, k2)
		}
	}

	return total
}

func inverseLevenshteinDistance(word1, word2 string) float64 {
	// Code from http://rosettacode.org/wiki/Levenshtein_distance#Go
	d := make([][]int, len(word1)+1)
	for i := range d {
		d[i] = make([]int, len(word2)+1)
	}
	for i := range d {
		d[i][0] = i
	}
	for j := range d[0] {
		d[0][j] = j
	}
	for j := 1; j <= len(word2); j++ {
		for i := 1; i <= len(word1); i++ {
			if word1[i-1] == word2[j-1] {
				d[i][j] = d[i-1][j-1]
			} else {
				min := d[i-1][j]
				if d[i][j-1] < min {
					min = d[i][j-1]
				}
				if d[i-1][j-1] < min {
					min = d[i-1][j-1]
				}
				d[i][j] = min + 1
			}
		}

	}
	return 1 / float64(d[len(word1)][len(word2)])
}
<commit_msg>Fix the `computeScore` function when distance reaches +Inf<commit_after>package player

import "math"

func computeScore(keywords1, keywords2 []string) float64 {
	var total float64 = 0

	for _, k1 := range keywords1 {
		for _, k2 := range keywords2 {
			d := inverseLevenshteinDistance(k1, k2)
			if math.IsInf(d, 0) {
				total = total + 10
			} else {
				total = total + d
			}
		}
	}

	return total
}

func inverseLevenshteinDistance(word1, word2 string) float64 {
	// Code from http://rosettacode.org/wiki/Levenshtein_distance#Go
	d := make([][]int, len(word1)+1)
	for i := range d {
		d[i] = make([]int, len(word2)+1)
	}
	for i := range d {
		d[i][0] = i
	}
	for j := range d[0] {
		d[0][j] = j
	}
	for j := 1; j <= len(word2); j++ {
		for i := 1; i <= len(word1); i++ {
			if word1[i-1] == word2[j-1] {
				d[i][j] = d[i-1][j-1]
			} else {
				min := d[i-1][j]
				if d[i][j-1] < min {
					min = d[i][j-1]
				}
				if d[i-1][j-1] < min {
					min = d[i-1][j-1]
				}
				d[i][j] = min + 1
			}
		}

	}
	return 1 / float64(d[len(word1)][len(word2)])
}
<|endoftext|>


<commit_before>package channels

type IdentifyDevice interface {
}

type IdentifyChannel struct {
	baseChannel
	device IdentifyDevice
}

func NewIdentifyChannel(device IdentifyDevice) *IdentifyChannel {
	return &IdentifyChannel{baseChannel{
		protocol: "identify",
	}, device}
}

func (c *IdentifyChannel) SendState(state float64) error {
	return c.SendEvent("state", state)
}
<commit_msg>add proper identify method to identify channel<commit_after>package channels

type IdentifyDevice interface {
	Identify() error
}

type IdentifyChannel struct {
	baseChannel
	device IdentifyDevice
}

func NewIdentifyChannel(device IdentifyDevice) *IdentifyChannel {
	return &IdentifyChannel{baseChannel{
		protocol: "identify",
	}, device}
}

func (c *IdentifyChannel) SendState(state float64) error {
	return c.SendEvent("state", state)
}

func (c *IdentifyChannel) Identify() error {
	return c.device.Identify()
}
<|endoftext|>


<commit_before>package archive

import (
	"archive/tar"
	"archive/zip"
	"compress/bzip2"
	"compress/gzip"
	"crypto/md5"
	"errors"
	"fmt"
	"github.com/MG-RAST/Shock/shock-server/conf"
	"io"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
)

var validUncompress = []string{"gzip", "bzip2"}
var validCompress = []string{"gzip", "zip"}
var validArchive = []string{"zip", "tar", "tar.gz", "tar.bz2"}
var ArchiveList = strings.Join(validArchive, ", ")

type FormFile struct {
	Name     string
	Path     string
	Checksum map[string]string
}

func IsValidArchive(a string) bool {
	for _, b := range validArchive {
		if b == a {
			return true
		}
	}
	return false
}

func IsValidUncompress(a string) bool {
	for _, b := range validUncompress {
		if b == a {
			return true
		}
	}
	return false
}

func IsValidCompress(a string) bool {
	for _, b := range validCompress {
		if b == a {
			return true
		}
	}
	return false
}

func FilesFromArchive(format string, filePath string) (fileList []FormFile, unpackDir string, err error) {
	// set unpack dir
	unpackDir = fmt.Sprintf("%s/temp/%d%d", conf.PATH_DATA, rand.Int(), rand.Int())
	if err = os.Mkdir(unpackDir, 0777); err != nil {
		return
	}
	// magic to unpack archive
	if format == "zip" {
		fileList, err = unZip(filePath, unpackDir)
	} else if format == "tar" {
		fileList, err = unTar(filePath, unpackDir, "")
	} else if format == "tar.gz" {
		fileList, err = unTar(filePath, unpackDir, "gzip")
	} else if format == "tar.bz2" {
		fileList, err = unTar(filePath, unpackDir, "bzip2")
	} else {
		return nil, unpackDir, errors.New("invalid archive format. must be one of: " + ArchiveList)
	}
	return
}

func unTar(filePath string, unpackDir string, compression string) (fileList []FormFile, err error) {
	// get file handle
	openFile, err := os.Open(filePath)
	if err != nil {
		return
	}
	defer openFile.Close()

	// tarball handle
	var tarBallReader *tar.Reader

	// filter through compression if use
	if compression == "gzip" {
		gReader, gerr := gzip.NewReader(openFile)
		if gerr != nil {
			return nil, gerr
		}
		defer gReader.Close()
		tarBallReader = tar.NewReader(gReader)
	} else if compression == "bzip2" {
		bReader := bzip2.NewReader(openFile)
		tarBallReader = tar.NewReader(bReader)
	} else {
		tarBallReader = tar.NewReader(openFile)
	}

	// extract tarball
	for {
		header, nerr := tarBallReader.Next()
		if nerr != nil {
			if nerr == io.EOF {
				break
			}
			return nil, nerr
		}

		// set names
		path := filepath.Join(unpackDir, header.Name)
		baseName := filepath.Base(header.Name)
		// skip hidden
		if strings.HasPrefix(baseName, ".") {
			continue
		}

		// only handle real files and dirs, ignore links
		switch header.Typeflag {
		// handle directory
		case tar.TypeDir:
			if err = os.MkdirAll(path, 0777); err != nil {
				return
			}
		// handle regualar file
		case tar.TypeReg:
			// open output file
			var writer *os.File
			if writer, err = os.Create(path); err != nil {
				return
			}
			// get md5
			md5h := md5.New()
			dst := io.MultiWriter(writer, md5h)
			// write it
			_, err = io.Copy(dst, tarBallReader)
			writer.Close()
			if err != nil {
				return
			}
			// add to filelist
			ffile := FormFile{Name: baseName, Path: path, Checksum: make(map[string]string)}
			ffile.Checksum["md5"] = fmt.Sprintf("%x", md5h.Sum(nil))
			fileList = append(fileList, ffile)
		default:
		}
	}
	return
}

func unZip(filePath string, unpackDir string) (fileList []FormFile, err error) {
	// open file with unzip
	zipReader, err := zip.OpenReader(filePath)
	if err != nil {
		return
	}

	// extract archive
	for _, zf := range zipReader.File {
		// set names
		path := filepath.Join(unpackDir, zf.Name)
		baseName := filepath.Base(zf.Name)
		// skip hidden
		if strings.HasPrefix(baseName, ".") {
			continue
		}

		if zf.FileInfo().IsDir() {
			// handle directory
			if err = os.MkdirAll(path, 0777); err != nil {
				return
			}
		} else {
			// open output file
			writer, werr := os.Create(path)
			if werr != nil {
				return nil, werr
			}
			// open input stream
			zfh, zerr := zf.Open()
			if zerr != nil {
				writer.Close()
				return nil, zerr
			}
			// get md5
			md5h := md5.New()
			dst := io.MultiWriter(writer, md5h)
			// write it
			_, err = io.Copy(dst, zfh)
			writer.Close()
			zfh.Close()
			if err != nil {
				return
			}
			// add to filelist
			ffile := FormFile{Name: baseName, Path: path, Checksum: make(map[string]string)}
			ffile.Checksum["md5"] = fmt.Sprintf("%x", md5h.Sum(nil))
			fileList = append(fileList, ffile)
		}
	}
	return
}

func CompressReader(format string, filename string, inReader io.ReadCloser) (outReader io.ReadCloser) {
	if IsValidCompress(format) {
		pReader, pWriter := io.Pipe()
		if format == "gzip" {
			gWriter := gzip.NewWriter(pWriter)
			gWriter.Header.Name = filename
			go func() {
				io.Copy(gWriter, inReader)
				gWriter.Close()
				pWriter.Close()
			}()
		} else if format == "zip" {
			zWriter := zip.NewWriter(pWriter)
			zFile, _ := zWriter.Create(filename)
			go func() {
				io.Copy(zFile, inReader)
				zWriter.Close()
				pWriter.Close()
			}()
		}
		return pReader
	}
	// just return input reader if no valid compression
	return inReader
}

func UncompressReader(format string, inReader io.Reader) (outReader io.Reader, err error) {
	if IsValidUncompress(format) {
		if format == "gzip" {
			gWriter, gerr := gzip.NewReader(inReader)
			if gerr != nil {
				return nil, gerr
			}
			return gWriter, nil
		} else if format == "bzip2" {
			bWriter := bzip2.NewReader(inReader)
			return bWriter, nil
		}
	}
	return inReader, nil
}
<commit_msg>more archive cleanup<commit_after>package archive

import (
	"archive/tar"
	"archive/zip"
	"compress/bzip2"
	"compress/gzip"
	"crypto/md5"
	"errors"
	"fmt"
	"github.com/MG-RAST/Shock/shock-server/conf"
	"io"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
)

var validUncompress = []string{"gzip", "bzip2"}
var validCompress = []string{"gzip", "zip"}
var validArchive = []string{"zip", "tar", "tar.gz", "tar.bz2"}
var ArchiveList = strings.Join(validArchive, ", ")

type FormFile struct {
	Name     string
	Path     string
	Checksum map[string]string
}

func IsValidArchive(a string) bool {
	for _, b := range validArchive {
		if b == a {
			return true
		}
	}
	return false
}

func IsValidUncompress(a string) bool {
	for _, b := range validUncompress {
		if b == a {
			return true
		}
	}
	return false
}

func IsValidCompress(a string) bool {
	for _, b := range validCompress {
		if b == a {
			return true
		}
	}
	return false
}

func FilesFromArchive(format string, filePath string) (fileList []FormFile, unpackDir string, err error) {
	// set unpack dir
	unpackDir = fmt.Sprintf("%s/temp/%d%d", conf.PATH_DATA, rand.Int(), rand.Int())
	if err = os.Mkdir(unpackDir, 0777); err != nil {
		return
	}
	// magic to unpack archive
	if format == "zip" {
		fileList, err = unZip(filePath, unpackDir)
	} else if format == "tar" {
		fileList, err = unTar(filePath, unpackDir, "")
	} else if format == "tar.gz" {
		fileList, err = unTar(filePath, unpackDir, "gzip")
	} else if format == "tar.bz2" {
		fileList, err = unTar(filePath, unpackDir, "bzip2")
	} else {
		return nil, unpackDir, errors.New("invalid archive format. must be one of: " + ArchiveList)
	}
	return
}

func unTar(filePath string, unpackDir string, compression string) (fileList []FormFile, err error) {
	// get file handle
	openFile, err := os.Open(filePath)
	if err != nil {
		return
	}
	defer openFile.Close()

	// filter through compression if use
	ucReader, ucErr := UncompressReader(compression, openFile)
	if ucErr != nil {
		return nil, ucErr
	}

	// extract tarball
	tarBallReader := tar.NewReader(ucReader)
	for {
		header, nerr := tarBallReader.Next()
		if nerr != nil {
			if nerr == io.EOF {
				break
			}
			return nil, nerr
		}

		// set names
		path := filepath.Join(unpackDir, header.Name)
		baseName := filepath.Base(header.Name)
		// skip hidden
		if strings.HasPrefix(baseName, ".") {
			continue
		}

		// only handle real files and dirs, ignore links
		switch header.Typeflag {
		// handle directory
		case tar.TypeDir:
			if err = os.MkdirAll(path, 0777); err != nil {
				return
			}
		// handle regualar file
		case tar.TypeReg:
			// open output file
			var writer *os.File
			if writer, err = os.Create(path); err != nil {
				return
			}
			// get md5
			md5h := md5.New()
			dst := io.MultiWriter(writer, md5h)
			// write it
			_, err = io.Copy(dst, tarBallReader)
			writer.Close()
			if err != nil {
				return
			}
			// add to filelist
			ffile := FormFile{Name: baseName, Path: path, Checksum: make(map[string]string)}
			ffile.Checksum["md5"] = fmt.Sprintf("%x", md5h.Sum(nil))
			fileList = append(fileList, ffile)
		default:
		}
	}
	return
}

func unZip(filePath string, unpackDir string) (fileList []FormFile, err error) {
	// open file with unzip
	zipReader, err := zip.OpenReader(filePath)
	if err != nil {
		return
	}

	// extract archive
	for _, zf := range zipReader.File {
		// set names
		path := filepath.Join(unpackDir, zf.Name)
		baseName := filepath.Base(zf.Name)
		// skip hidden
		if strings.HasPrefix(baseName, ".") {
			continue
		}

		if zf.FileInfo().IsDir() {
			// handle directory
			if err = os.MkdirAll(path, 0777); err != nil {
				return
			}
		} else {
			// open output file
			writer, werr := os.Create(path)
			if werr != nil {
				return nil, werr
			}
			// open input stream
			zfh, zerr := zf.Open()
			if zerr != nil {
				writer.Close()
				return nil, zerr
			}
			// get md5
			md5h := md5.New()
			dst := io.MultiWriter(writer, md5h)
			// write it
			_, err = io.Copy(dst, zfh)
			writer.Close()
			zfh.Close()
			if err != nil {
				return
			}
			// add to filelist
			ffile := FormFile{Name: baseName, Path: path, Checksum: make(map[string]string)}
			ffile.Checksum["md5"] = fmt.Sprintf("%x", md5h.Sum(nil))
			fileList = append(fileList, ffile)
		}
	}
	return
}

func CompressReader(format string, filename string, inReader io.ReadCloser) (outReader io.ReadCloser) {
	if IsValidCompress(format) {
		pReader, pWriter := io.Pipe()
		if format == "gzip" {
			gWriter := gzip.NewWriter(pWriter)
			gWriter.Header.Name = filename
			go func() {
				io.Copy(gWriter, inReader)
				gWriter.Close()
				pWriter.Close()
			}()
		} else if format == "zip" {
			zWriter := zip.NewWriter(pWriter)
			zFile, _ := zWriter.Create(filename)
			go func() {
				io.Copy(zFile, inReader)
				zWriter.Close()
				pWriter.Close()
			}()
		}
		return pReader
	}
	// just return input reader if no valid compression
	return inReader
}

func UncompressReader(format string, inReader io.Reader) (outReader io.Reader, err error) {
	if IsValidUncompress(format) {
		if format == "gzip" {
			gWriter, gerr := gzip.NewReader(inReader)
			if gerr != nil {
				return nil, gerr
			}
			return gWriter, nil
		} else if format == "bzip2" {
			bWriter := bzip2.NewReader(inReader)
			return bWriter, nil
		}
	}
	return inReader, nil
}
<|endoftext|>


<commit_before>package http

import (
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/oinume/lekcije/server/context_data"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestIndex(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	r := require.New(t)
	err := os.Chdir("../../..")
	r.NoError(err)

	req, err := http.NewRequest("GET", "/", nil)
	r.NoError(err)
	ctx := context_data.SetTrackingID(req.Context(), "a")
	req = req.WithContext(ctx)

	rr := httptest.NewRecorder()
	http.HandlerFunc(Index).ServeHTTP(rr, req)
	a.Equal(http.StatusOK, rr.Code)
	a.Contains(rr.Body.String(), "<title>lekcije - DMM英会話のお気に入り講師をフォローしよう</title>")
}
<commit_msg>Use table driven test<commit_after>package http

import (
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/oinume/lekcije/server/context_data"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGet(t *testing.T) {
	t.Parallel()
	a := assert.New(t)
	r := require.New(t)
	err := os.Chdir("../../..")
	r.NoError(err)

	testCases := []struct {
		path     string
		handler  http.HandlerFunc
		code     int
		keywords []string
	}{
		{
			path:     "/",
			handler:  http.HandlerFunc(Index),
			code:     http.StatusOK,
			keywords: []string{"<title>lekcije - DMM英会話のお気に入り講師をフォローしよう</title>"},
		},
		{
			path:     "/robots.txt",
			handler:  http.HandlerFunc(RobotsTxt),
			code:     http.StatusOK,
			keywords: []string{"Allow: /"},
		},
	}

	for _, tc := range testCases {
		req, err := http.NewRequest("GET", tc.path, nil)
		r.NoError(err)

		ctx := context_data.SetTrackingID(req.Context(), "a")
		req = req.WithContext(ctx)

		rr := httptest.NewRecorder()
		tc.handler.ServeHTTP(rr, req)
		a.Equal(tc.code, rr.Code)
		for _, keyword := range tc.keywords {
			a.Contains(rr.Body.String(), keyword)
		}
	}
}
<|endoftext|>


<commit_before>package codeflow

import (
	"log"
	"net/http"
	"time"

	"gopkg.in/mgo.v2/bson"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/checkr/codeflow/server/plugins"
)

type Admin struct {
	Path string
}

func (x *Admin) Register(api *rest.Api) []*rest.Route {
	var routes []*rest.Route
	routes = append(routes,
		rest.Get(x.Path+"/serviceSpecs", x.serviceSpecs),
		rest.Put(x.Path+"/serviceSpecs", x.updateServiceSpec),
		rest.Delete(x.Path+"/serviceSpecs/#slug", x.deleteServiceSpec),
		rest.Get(x.Path+"/serviceSpecs/#slug/services", x.serviceSpecServices),
		rest.Put(x.Path+"/deployallprojects", x.deployAllProjects),
		rest.Put(x.Path+"/deployloadbalancers", x.deployAllLoadBalancers),
	)
	log.Printf("Started the codeflow admin handler on %s\n", x.Path)
	return routes
}

// Deploys all projects (does not touch services)
func (x *Admin) deployAllProjects(w rest.ResponseWriter, r *rest.Request) {
	projects := []Project{}
	m := bson.M{"deleted": false}
	results := db.Collection("projects").Find(m)
	results.Query.Sort("-$natural").All(&projects)
	log.Printf("Deploying %d projects", len(projects))
	for _, p := range projects {
		log.Printf("Deploying %s...", p.Slug)
		release := Release{}
		if err := GetCurrentRelease(p.Id, &release); err != nil {
			log.Printf("GetCurrentRelease::Error: %s; Aborting release for %s", err, p.Slug)
		} else {
			log.Printf("Deploying %s release", release.Id)
			err := CreateDeploy(&release)
			if err != nil {
				log.Printf("ERROR '%s' creating deployment for %s", err, p.Slug)
			}
		}
	}
}

// Deploys all loadBalancers
func (x *Admin) deployAllLoadBalancers(w rest.ResponseWriter, r *rest.Request) {
	projects := []Project{}
	m := bson.M{"deleted": false}
	results := db.Collection("projects").Find(m)
	results.Query.Sort("-$natural").All(&projects)
	for _, p := range projects {
		log.Printf("Finding services for '%s' project...", p.Slug)
		l := bson.M{"projectId": p.Id, "state": "complete"}
		lbResults := db.Collection("extensions").Find(l)
		loadBalancer := &LoadBalancer{}
		for lbResults.Next(loadBalancer) {
			log.Printf("Deploying loadbalancer %s for %s", loadBalancer.Name, p.Slug)
			err := ExtensionUpdated(loadBalancer)
			if err != nil {
				log.Printf("ERROR '%s' updating loadBalancer %s", err, loadBalancer.Name)
			}
		}
	}
	// Throttle so that we can be "nice" to the amazon APIs
	log.Println("Sleeping 1 second")
	time.Sleep(time.Second * 1)
}

func (x *Admin) deleteServiceSpec(w rest.ResponseWriter, r *rest.Request) {
	slug := r.PathParam("slug")

	match := bson.M{"_id": bson.ObjectIdHex(slug)}

	if err := db.Collection("serviceSpecs").DeleteOne(match); err != nil {
		log.Printf("ServiceSpec:: Delete Error %s", err)
		return
	}
}

func (x *Admin) updateServiceSpec(w rest.ResponseWriter, r *rest.Request) {
	spec := ServiceSpec{}

	err := r.DecodeJsonPayload(&spec)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if err := db.Collection("serviceSpecs").Save(&spec); err != nil {
		log.Printf("ServiceSpec:: Save Error %s", err)
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteJson(&spec)
}

func (x *Admin) serviceSpecs(w rest.ResponseWriter, r *rest.Request) {

	specs := []ServiceSpec{}
	spec := ServiceSpec{}

	results := db.Collection("serviceSpecs").Find(bson.M{})
	for results.Next(&spec) {
		specs = append(specs, spec)
	}

	w.WriteJson(specs)
}

func (x *Admin) serviceSpecServices(w rest.ResponseWriter, r *rest.Request) {
	services := []Service{}
	service := Service{}

	slug := r.PathParam("slug")
	match := bson.M{"specId": bson.ObjectIdHex(slug), "state": bson.M{"$in": []plugins.State{plugins.Waiting, plugins.Running}}}

	results := db.Collection("services").Find(match)
	for results.Next(&service) {
		services = append(services, service)
	}

	w.WriteJson(services)
}
<commit_msg>exclude checkr-codeflow from deployall* (#164)<commit_after>package codeflow

import (
	"log"
	"net/http"
	"time"

	"gopkg.in/mgo.v2/bson"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/checkr/codeflow/server/plugins"
)

type Admin struct {
	Path string
}

func (x *Admin) Register(api *rest.Api) []*rest.Route {
	var routes []*rest.Route
	routes = append(routes,
		rest.Get(x.Path+"/serviceSpecs", x.serviceSpecs),
		rest.Put(x.Path+"/serviceSpecs", x.updateServiceSpec),
		rest.Delete(x.Path+"/serviceSpecs/#slug", x.deleteServiceSpec),
		rest.Get(x.Path+"/serviceSpecs/#slug/services", x.serviceSpecServices),
		rest.Put(x.Path+"/deployallprojects", x.deployAllProjects),
		rest.Put(x.Path+"/deployloadbalancers", x.deployAllLoadBalancers),
	)
	log.Printf("Started the codeflow admin handler on %s\n", x.Path)
	return routes
}

// Deploys all projects (does not touch services)
func (x *Admin) deployAllProjects(w rest.ResponseWriter, r *rest.Request) {
	projects := []Project{}
	m := bson.M{"deleted": false, "slug": bson.M{"$ne": "checkr-codeflow"}}
	results := db.Collection("projects").Find(m)
	results.Query.Sort("-$natural").All(&projects)
	log.Printf("Deploying %d projects", len(projects))
	for _, p := range projects {
		log.Printf("Deploying %s...", p.Slug)
		release := Release{}
		if err := GetCurrentRelease(p.Id, &release); err != nil {
			log.Printf("GetCurrentRelease::Error: %s; Aborting release for %s", err, p.Slug)
		} else {
			log.Printf("Deploying %s release", release.Id)
			err := CreateDeploy(&release)
			if err != nil {
				log.Printf("ERROR '%s' creating deployment for %s", err, p.Slug)
			}
		}
	}
}

// Deploys all loadBalancers
func (x *Admin) deployAllLoadBalancers(w rest.ResponseWriter, r *rest.Request) {
	projects := []Project{}
	m := bson.M{"deleted": false, "slug": bson.M{"$ne": "checkr-codeflow"}}
	results := db.Collection("projects").Find(m)
	results.Query.Sort("-$natural").All(&projects)
	for _, p := range projects {
		log.Printf("Finding services for '%s' project...", p.Slug)
		l := bson.M{"projectId": p.Id, "state": "complete"}
		lbResults := db.Collection("extensions").Find(l)
		loadBalancer := &LoadBalancer{}
		for lbResults.Next(loadBalancer) {
			log.Printf("Deploying loadbalancer %s for %s", loadBalancer.Name, p.Slug)
			err := ExtensionUpdated(loadBalancer)
			if err != nil {
				log.Printf("ERROR '%s' updating loadBalancer %s", err, loadBalancer.Name)
			}
		}
	}
	// Throttle so that we can be "nice" to the amazon APIs
	log.Println("Sleeping 1 second")
	time.Sleep(time.Second * 1)
}

func (x *Admin) deleteServiceSpec(w rest.ResponseWriter, r *rest.Request) {
	slug := r.PathParam("slug")

	match := bson.M{"_id": bson.ObjectIdHex(slug)}

	if err := db.Collection("serviceSpecs").DeleteOne(match); err != nil {
		log.Printf("ServiceSpec:: Delete Error %s", err)
		return
	}
}

func (x *Admin) updateServiceSpec(w rest.ResponseWriter, r *rest.Request) {
	spec := ServiceSpec{}

	err := r.DecodeJsonPayload(&spec)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if err := db.Collection("serviceSpecs").Save(&spec); err != nil {
		log.Printf("ServiceSpec:: Save Error %s", err)
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteJson(&spec)
}

func (x *Admin) serviceSpecs(w rest.ResponseWriter, r *rest.Request) {

	specs := []ServiceSpec{}
	spec := ServiceSpec{}

	results := db.Collection("serviceSpecs").Find(bson.M{})
	for results.Next(&spec) {
		specs = append(specs, spec)
	}

	w.WriteJson(specs)
}

func (x *Admin) serviceSpecServices(w rest.ResponseWriter, r *rest.Request) {
	services := []Service{}
	service := Service{}

	slug := r.PathParam("slug")
	match := bson.M{"specId": bson.ObjectIdHex(slug), "state": bson.M{"$in": []plugins.State{plugins.Waiting, plugins.Running}}}

	results := db.Collection("services").Find(match)
	for results.Next(&service) {
		services = append(services, service)
	}

	w.WriteJson(services)
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2017 The btcsuite developers
// Copyright (c) 2015-2016 The Decred developers
// Heavily inspired by https://github.com/btcsuite/btcd/blob/master/version.go
// Copyright (C) 2015-2017 The Lightning Network Developers

package build

import (
	"bytes"
	"fmt"
	"strings"
)

// Commit stores the current commit hash of this build, this should be set using
// the -ldflags during compilation.
var Commit string

// semanticAlphabet is the set of characters that are permitted for use in an
// AppPreRelease.
const semanticAlphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-"

// These constants define the application version and follow the semantic
// versioning 2.0.0 spec (http://semver.org/).
const (
	appMajor uint = 0
	appMinor uint = 9
	appPatch uint = 0

	// appPreRelease MUST only contain characters from semanticAlphabet
	// per the semantic versioning spec.
	appPreRelease = "beta"
)

// Version returns the application version as a properly formed string per the
// semantic versioning 2.0.0 spec (http://semver.org/).
func Version() string {
	// Start with the major, minor, and patch versions.
	version := fmt.Sprintf("%d.%d.%d", appMajor, appMinor, appPatch)

	// Append pre-release version if there is one.  The hyphen called for
	// by the semantic versioning spec is automatically appended and should
	// not be contained in the pre-release string.  The pre-release version
	// is not appended if it contains invalid characters.
	preRelease := normalizeVerString(appPreRelease)
	if preRelease != "" {
		version = fmt.Sprintf("%s-%s", version, preRelease)
	}

	// Append commit hash of current build to version.
	version = fmt.Sprintf("%s commit=%s", version, Commit)

	return version
}

// normalizeVerString returns the passed string stripped of all characters which
// are not valid according to the semantic versioning guidelines for pre-release
// version and build metadata strings.  In particular they MUST only contain
// characters in semanticAlphabet.
func normalizeVerString(str string) string {
	var result bytes.Buffer
	for _, r := range str {
		if strings.ContainsRune(semanticAlphabet, r) {
			result.WriteRune(r)
		}
	}
	return result.String()
}
<commit_msg>build/version: expose AppMajor, AppMinor, AppPatch, etc.<commit_after>// Copyright (c) 2013-2017 The btcsuite developers
// Copyright (c) 2015-2016 The Decred developers
// Heavily inspired by https://github.com/btcsuite/btcd/blob/master/version.go
// Copyright (C) 2015-2017 The Lightning Network Developers

package build

import (
	"bytes"
	"fmt"
	"strings"
)

// Commit stores the current commit hash of this build, this should be set using
// the -ldflags during compilation.
var Commit string

// semanticAlphabet is the set of characters that are permitted for use in an
// AppPreRelease.
const semanticAlphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-"

// These constants define the application version and follow the semantic
// versioning 2.0.0 spec (http://semver.org/).
const (
	// AppMajor defines the major version of this binary.
	AppMajor uint = 0

	// AppMinor defines the minor version of this binary.
	AppMinor uint = 9

	// AppPatch defines the application patch for this binary.
	AppPatch uint = 0

	// AppPreRelease MUST only contain characters from semanticAlphabet
	// per the semantic versioning spec.
	AppPreRelease = "beta"
)

// Version returns the application version as a properly formed string per the
// semantic versioning 2.0.0 spec (http://semver.org/).
func Version() string {
	// Start with the major, minor, and patch versions.
	version := fmt.Sprintf("%d.%d.%d", AppMajor, AppMinor, AppPatch)

	// Append pre-release version if there is one.  The hyphen called for
	// by the semantic versioning spec is automatically appended and should
	// not be contained in the pre-release string.  The pre-release version
	// is not appended if it contains invalid characters.
	preRelease := normalizeVerString(AppPreRelease)
	if preRelease != "" {
		version = fmt.Sprintf("%s-%s", version, preRelease)
	}

	// Append commit hash of current build to version.
	version = fmt.Sprintf("%s commit=%s", version, Commit)

	return version
}

// normalizeVerString returns the passed string stripped of all characters which
// are not valid according to the semantic versioning guidelines for pre-release
// version and build metadata strings.  In particular they MUST only contain
// characters in semanticAlphabet.
func normalizeVerString(str string) string {
	var result bytes.Buffer
	for _, r := range str {
		if strings.ContainsRune(semanticAlphabet, r) {
			result.WriteRune(r)
		}
	}
	return result.String()
}
<|endoftext|>


<commit_before>package builder

// nginx
const (
	NginxVersion           = "1.21.6"
	NginxDownloadURLPrefix = "https://nginx.org/download"
)

// pcre
const (
	PcreVersion           = "10.39"
	PcreDownloadURLPrefix = "https://github.com/PhilipHazel/pcre2/releases/download"
)

// openssl
const (
	OpenSSLVersion           = "1.1.1n"
	OpenSSLDownloadURLPrefix = "https://www.openssl.org/source"
)

// libressl
const (
	LibreSSLVersion           = "3.4.2"
	LibreSSLDownloadURLPrefix = "https://ftp.openbsd.org/pub/OpenBSD/LibreSSL"
)

// zlib
const (
	ZlibVersion           = "1.2.11"
	ZlibDownloadURLPrefix = "https://zlib.net/fossils"
)

// openResty
const (
	OpenRestyVersion           = "1.19.9.1"
	OpenRestyDownloadURLPrefix = "https://openresty.org/download"
)

// tengine
const (
	TengineVersion           = "2.3.3"
	TengineDownloadURLPrefix = "https://tengine.taobao.org/download"
)

// component enumerations
const (
	ComponentNginx = iota
	ComponentOpenResty
	ComponentTengine
	ComponentPcre
	ComponentOpenSSL
	ComponentLibreSSL
	ComponentZlib
	ComponentMax
)
<commit_msg>bumped default zlib version to 1.2.12.<commit_after>package builder

// nginx
const (
	NginxVersion           = "1.21.6"
	NginxDownloadURLPrefix = "https://nginx.org/download"
)

// pcre
const (
	PcreVersion           = "10.39"
	PcreDownloadURLPrefix = "https://github.com/PhilipHazel/pcre2/releases/download"
)

// openssl
const (
	OpenSSLVersion           = "1.1.1n"
	OpenSSLDownloadURLPrefix = "https://www.openssl.org/source"
)

// libressl
const (
	LibreSSLVersion           = "3.4.2"
	LibreSSLDownloadURLPrefix = "https://ftp.openbsd.org/pub/OpenBSD/LibreSSL"
)

// zlib
const (
	ZlibVersion           = "1.2.12"
	ZlibDownloadURLPrefix = "https://zlib.net"
)

// openResty
const (
	OpenRestyVersion           = "1.19.9.1"
	OpenRestyDownloadURLPrefix = "https://openresty.org/download"
)

// tengine
const (
	TengineVersion           = "2.3.3"
	TengineDownloadURLPrefix = "https://tengine.taobao.org/download"
)

// component enumerations
const (
	ComponentNginx = iota
	ComponentOpenResty
	ComponentTengine
	ComponentPcre
	ComponentOpenSSL
	ComponentLibreSSL
	ComponentZlib
	ComponentMax
)
<|endoftext|>


<commit_before>package builder

// nginx
const (
	NginxVersion           = "1.11.4"
	NginxDownloadURLPrefix = "http://nginx.org/download"
)

// pcre
const (
	PcreVersion           = "8.39"
	PcreDownloadURLPrefix = "http://ftp.csx.cam.ac.uk/pub/software/programming/pcre"
)

// openssl
const (
	OpenSSLVersion           = "1.0.2h"
	OpenSSLDownloadURLPrefix = "http://www.openssl.org/source"
)

// zlib
const (
	ZlibVersion           = "1.2.8"
	ZlibDownloadURLPrefix = "http://zlib.net"
)

// openResty
const (
	OpenRestyVersion           = "1.11.2.1"
	OpenRestyDownloadURLPrefix = "https://openresty.org/download"
)

// tengine
const (
	TengineVersion           = "2.1.2"
	TengineDownloadURLPrefix = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	ComponentNginx = iota
	ComponentOpenResty
	ComponentTengine
	ComponentPcre
	ComponentOpenSSL
	ComponentZlib
	ComponentMax
)
<commit_msg>bumped OpenSSL version to 1.0.2i by default.<commit_after>package builder

// nginx
const (
	NginxVersion           = "1.11.4"
	NginxDownloadURLPrefix = "http://nginx.org/download"
)

// pcre
const (
	PcreVersion           = "8.39"
	PcreDownloadURLPrefix = "http://ftp.csx.cam.ac.uk/pub/software/programming/pcre"
)

// openssl
const (
	OpenSSLVersion           = "1.0.2i"
	OpenSSLDownloadURLPrefix = "http://www.openssl.org/source"
)

// zlib
const (
	ZlibVersion           = "1.2.8"
	ZlibDownloadURLPrefix = "http://zlib.net"
)

// openResty
const (
	OpenRestyVersion           = "1.11.2.1"
	OpenRestyDownloadURLPrefix = "https://openresty.org/download"
)

// tengine
const (
	TengineVersion           = "2.1.2"
	TengineDownloadURLPrefix = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	ComponentNginx = iota
	ComponentOpenResty
	ComponentTengine
	ComponentPcre
	ComponentOpenSSL
	ComponentZlib
	ComponentMax
)
<|endoftext|>


<commit_before>package servicediscovery

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/hyperhq/hyper/utils"
	"github.com/hyperhq/runv/hypervisor"
	"github.com/hyperhq/runv/hypervisor/pod"
	"github.com/hyperhq/runv/lib/glog"
)

var (
	ServiceVolume string = "/usr/local/etc/haproxy/"
	ServiceImage  string = "haproxy:latest"
	ServiceConfig string = "haproxy.cfg"
)

func ApplyServices(vm *hypervisor.Vm, container string, services []pod.UserService) error {
	config := path.Join(ServiceVolume, ServiceConfig)

	vm.WriteFile(container, config, GenerateServiceConfig(services))

	command := "haproxy -D -f /usr/local/etc/haproxy/haproxy.cfg -p /var/run/haproxy.pid -sf `cat /var/run/haproxy.pid`"
	execcmd, err := json.Marshal(command)
	if err != nil {
		return err
	}

	return vm.Exec(nil, nil, string(execcmd), "", container)
}

func GetServices(vm *hypervisor.Vm, container string) ([]pod.UserService, error) {
	var services []pod.UserService
	config := path.Join(ServiceVolume, ServiceConfig)

	data, err := vm.ReadFile(container, config)
	if err != nil {
		return nil, err
	}

	token := bytes.Split(data, []byte("\n"))

	for _, tok := range token {
		first := bytes.Split(tok, []byte(" "))
		reader := bytes.NewReader(tok)
		if len(first) > 0 {
			var t1, t2, t3, t4 string
			if string(first[0][:]) == "frontend" {
				s := pod.UserService{
					Protocol: "TCP",
				}

				_, err := fmt.Fscanf(reader, "%s %s %s", &t1, &t2, &t3)
				if err != nil {
					return nil, err
				}

				hostport := strings.Split(t3, ":")
				s.ServiceIP = hostport[0]
				port, err := strconv.ParseInt(hostport[1], 10, 32)
				if err != nil {
					return nil, err
				}
				s.ServicePort = int(port)

				services = append(services, s)
			} else if string(first[0][:]) == "\tserver" {
				var idx int
				var h pod.UserServiceBackend
				_, err := fmt.Fscanf(reader, "%s %s %s %s", &t1, &t2, &t3, &t4)
				if err != nil {
					return nil, err
				}

				hostport := strings.Split(t3, ":")
				h.HostIP = hostport[0]
				port, err := strconv.ParseInt(hostport[1], 10, 32)
				if err != nil {
					return nil, err
				}
				h.HostPort = int(port)

				idxs := strings.Split(t2, "-")
				idxLong, err := strconv.ParseInt(idxs[1], 10, 32)
				if err != nil {
					return nil, err
				}
				idx = int(idxLong)

				services[idx].Hosts = append(services[idx].Hosts, h)
			}
		}
	}
	return services, nil
}

func GenerateServiceConfig(services []pod.UserService) []byte {
	data := []byte{}

	globalConfig := fmt.Sprintf("global\n\t#chroot\t/var/lib/haproxy\n\tpidfile\t/var/run/haproxy.pid\n\tmaxconn\t4000\n\t#user\thaproxy\n\t#group\thaproxy\n\tdaemon\ndefaults\n\tmode\ttcp\n\tretries\t3\n\ttimeout queue\t1m\n\ttimeout connect\t10s\n\ttimeout client\t1m\n\ttimeout server\t1m\n\ttimeout check\t10s\n\tmaxconn\t3000\n")

	data = append(data, globalConfig...)
	for idx, srv := range services {
		front := fmt.Sprintf("frontend front%d %s:%d\n\tdefault_backend\tback%d\n",
			idx, srv.ServiceIP, srv.ServicePort, idx)
		data = append(data, front...)
		back := fmt.Sprintf("backend back%d\n\tbalance\troundrobin\n", idx)
		data = append(data, back...)
		for hostid, host := range srv.Hosts {
			back := fmt.Sprintf("\tserver back-%d-%d %s:%d check\n",
				idx, hostid, host.HostIP, host.HostPort)
			data = append(data, back...)
		}
	}

	glog.V(1).Infof("haproxy config: %s", data[:])
	return data
}

func checkHaproxyConfig(services []pod.UserService, config string) error {
	var err error
	glog.V(1).Infof("haproxy config: %s\n", config)
	if _, err = os.Stat(config); err != nil && os.IsNotExist(err) {
		/* Generate haproxy config from service and write to config */
		return ioutil.WriteFile(config, GenerateServiceConfig(services), 0644)
	}
	return err
}

func PrepareServices(userPod *pod.UserPod, podId string) error {
	var serviceDir string = path.Join(utils.HYPER_ROOT, "services", podId)
	var config string = path.Join(serviceDir, ServiceConfig)
	var err error

	if len(userPod.Services) == 0 {
		return nil
	}

	if err = os.MkdirAll(serviceDir, 0755); err != nil && !os.IsExist(err) {
		return err
	}

	return checkHaproxyConfig(userPod.Services, config)
}
<commit_msg>Fix vm.Exec command error<commit_after>package servicediscovery

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/hyperhq/hyper/utils"
	"github.com/hyperhq/runv/hypervisor"
	"github.com/hyperhq/runv/hypervisor/pod"
	"github.com/hyperhq/runv/lib/glog"
)

var (
	ServiceVolume string = "/usr/local/etc/haproxy/"
	ServiceImage  string = "haproxy:latest"
	ServiceConfig string = "haproxy.cfg"
)

func ApplyServices(vm *hypervisor.Vm, container string, services []pod.UserService) error {
	config := path.Join(ServiceVolume, ServiceConfig)

	vm.WriteFile(container, config, GenerateServiceConfig(services))

	command := "haproxy -D -f /usr/local/etc/haproxy/haproxy.cfg -p /var/run/haproxy.pid -sf `cat /var/run/haproxy.pid`"
	execcmd, err := json.Marshal(strings.Split(command, " "))
	if err != nil {
		return err
	}

	return vm.Exec(nil, nil, string(execcmd), "", container)
}

func GetServices(vm *hypervisor.Vm, container string) ([]pod.UserService, error) {
	var services []pod.UserService
	config := path.Join(ServiceVolume, ServiceConfig)

	data, err := vm.ReadFile(container, config)
	if err != nil {
		return nil, err
	}

	token := bytes.Split(data, []byte("\n"))

	for _, tok := range token {
		first := bytes.Split(tok, []byte(" "))
		reader := bytes.NewReader(tok)
		if len(first) > 0 {
			var t1, t2, t3, t4 string
			if string(first[0][:]) == "frontend" {
				s := pod.UserService{
					Protocol: "TCP",
				}

				_, err := fmt.Fscanf(reader, "%s %s %s", &t1, &t2, &t3)
				if err != nil {
					return nil, err
				}

				hostport := strings.Split(t3, ":")
				s.ServiceIP = hostport[0]
				port, err := strconv.ParseInt(hostport[1], 10, 32)
				if err != nil {
					return nil, err
				}
				s.ServicePort = int(port)

				services = append(services, s)
			} else if string(first[0][:]) == "\tserver" {
				var idx int
				var h pod.UserServiceBackend
				_, err := fmt.Fscanf(reader, "%s %s %s %s", &t1, &t2, &t3, &t4)
				if err != nil {
					return nil, err
				}

				hostport := strings.Split(t3, ":")
				h.HostIP = hostport[0]
				port, err := strconv.ParseInt(hostport[1], 10, 32)
				if err != nil {
					return nil, err
				}
				h.HostPort = int(port)

				idxs := strings.Split(t2, "-")
				idxLong, err := strconv.ParseInt(idxs[1], 10, 32)
				if err != nil {
					return nil, err
				}
				idx = int(idxLong)

				services[idx].Hosts = append(services[idx].Hosts, h)
			}
		}
	}
	return services, nil
}

func GenerateServiceConfig(services []pod.UserService) []byte {
	data := []byte{}

	globalConfig := fmt.Sprintf("global\n\t#chroot\t/var/lib/haproxy\n\tpidfile\t/var/run/haproxy.pid\n\tmaxconn\t4000\n\t#user\thaproxy\n\t#group\thaproxy\n\tdaemon\ndefaults\n\tmode\ttcp\n\tretries\t3\n\ttimeout queue\t1m\n\ttimeout connect\t10s\n\ttimeout client\t1m\n\ttimeout server\t1m\n\ttimeout check\t10s\n\tmaxconn\t3000\n")

	data = append(data, globalConfig...)
	for idx, srv := range services {
		front := fmt.Sprintf("frontend front%d %s:%d\n\tdefault_backend\tback%d\n",
			idx, srv.ServiceIP, srv.ServicePort, idx)
		data = append(data, front...)
		back := fmt.Sprintf("backend back%d\n\tbalance\troundrobin\n", idx)
		data = append(data, back...)
		for hostid, host := range srv.Hosts {
			back := fmt.Sprintf("\tserver back-%d-%d %s:%d check\n",
				idx, hostid, host.HostIP, host.HostPort)
			data = append(data, back...)
		}
	}

	glog.V(1).Infof("haproxy config: %s", data[:])
	return data
}

func checkHaproxyConfig(services []pod.UserService, config string) error {
	var err error
	glog.V(1).Infof("haproxy config: %s\n", config)
	if _, err = os.Stat(config); err != nil && os.IsNotExist(err) {
		/* Generate haproxy config from service and write to config */
		return ioutil.WriteFile(config, GenerateServiceConfig(services), 0644)
	}
	return err
}

func PrepareServices(userPod *pod.UserPod, podId string) error {
	var serviceDir string = path.Join(utils.HYPER_ROOT, "services", podId)
	var config string = path.Join(serviceDir, ServiceConfig)
	var err error

	if len(userPod.Services) == 0 {
		return nil
	}

	if err = os.MkdirAll(serviceDir, 0755); err != nil && !os.IsExist(err) {
		return err
	}

	return checkHaproxyConfig(userPod.Services, config)
}
<|endoftext|>