<commit_before>package csrf

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
)

func init() {
	gin.SetMode(gin.TestMode)
}

func newServer(options Options) *gin.Engine {
	g := gin.New()
	store := sessions.NewCookieStore([]byte("secret123"))

	g.Use(sessions.Sessions("my_session", store))
	g.Use(Middleware(options))

	return g
}

type requestOptions struct {
	Method  string
	URL     string
	Headers map[string]string
	Body    io.Reader
}

func request(server *gin.Engine, options requestOptions) *httptest.ResponseRecorder {
	if options.Method == "" {
		options.Method = "GET"
	}

	w := httptest.NewRecorder()
	req, err := http.NewRequest(options.Method, options.URL, options.Body)

	if options.Headers != nil {
		for key, value := range options.Headers {
			req.Header.Set(key, value)
		}
	}

	server.ServeHTTP(w, req)

	if err != nil {
		panic(err)
	}

	return w
}

func TestForm(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"Content-Type": "application/x-www-form-urlencoded",
		},
		Body: strings.NewReader("_csrf=" + token),
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestQueryString(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login?_csrf=" + token,
		Headers: map[string]string{
			"Cookie": r1.Header().Get("Set-Cookie"),
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestQueryHeader1(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"X-CSRF-Token": token,
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestQueryHeader2(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"X-XSRF-Token": token,
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestErrorFunc(t *testing.T) {
	result := ""
	g := newServer(Options{
		Secret: "secret123",
		ErrorFunc: func(c *gin.Context) {
			result = "something wrong"
		},
	})

	g.GET("/login", func(c *gin.Context) {
		GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie": r1.Header().Get("Set-Cookie"),
		},
	})

	if result != "something wrong" {
		t.Error("Error function was not called")
	}
}

func TestIgnoreMethods(t *testing.T) {
	g := newServer(Options{
		Secret:        "secret123",
		IgnoreMethods: []string{"GET", "POST"},
	})

	g.GET("/login", func(c *gin.Context) {
		GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie": r1.Header().Get("Set-Cookie"),
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestTokenGetter(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
		TokenGetter: func(c *gin.Context) string {
			return c.Request.FormValue("wtf")
		},
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"Content-Type": "application/x-www-form-urlencoded",
		},
		Body: strings.NewReader("wtf=" + token),
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}
<commit_msg>fix test<commit_after>package csrf

import (
	"github.com/gin-contrib/sessions/cookie"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
)

func init() {
	gin.SetMode(gin.TestMode)
}

func newServer(options Options) *gin.Engine {
	g := gin.New()

	store := cookie.NewStore([]byte("secret123"))

	g.Use(sessions.Sessions("my_session", store))
	g.Use(Middleware(options))

	return g
}

type requestOptions struct {
	Method  string
	URL     string
	Headers map[string]string
	Body    io.Reader
}

func request(server *gin.Engine, options requestOptions) *httptest.ResponseRecorder {
	if options.Method == "" {
		options.Method = "GET"
	}

	w := httptest.NewRecorder()
	req, err := http.NewRequest(options.Method, options.URL, options.Body)

	if options.Headers != nil {
		for key, value := range options.Headers {
			req.Header.Set(key, value)
		}
	}

	server.ServeHTTP(w, req)

	if err != nil {
		panic(err)
	}

	return w
}

func TestForm(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"Content-Type": "application/x-www-form-urlencoded",
		},
		Body: strings.NewReader("_csrf=" + token),
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestQueryString(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login?_csrf=" + token,
		Headers: map[string]string{
			"Cookie": r1.Header().Get("Set-Cookie"),
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestQueryHeader1(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"X-CSRF-Token": token,
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestQueryHeader2(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"X-XSRF-Token": token,
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestErrorFunc(t *testing.T) {
	result := ""
	g := newServer(Options{
		Secret: "secret123",
		ErrorFunc: func(c *gin.Context) {
			result = "something wrong"
		},
	})

	g.GET("/login", func(c *gin.Context) {
		GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie": r1.Header().Get("Set-Cookie"),
		},
	})

	if result != "something wrong" {
		t.Error("Error function was not called")
	}
}

func TestIgnoreMethods(t *testing.T) {
	g := newServer(Options{
		Secret:        "secret123",
		IgnoreMethods: []string{"GET", "POST"},
	})

	g.GET("/login", func(c *gin.Context) {
		GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie": r1.Header().Get("Set-Cookie"),
		},
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}

func TestTokenGetter(t *testing.T) {
	var token string
	g := newServer(Options{
		Secret: "secret123",
		TokenGetter: func(c *gin.Context) string {
			return c.Request.FormValue("wtf")
		},
	})

	g.GET("/login", func(c *gin.Context) {
		token = GetToken(c)
	})

	g.POST("/login", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	r1 := request(g, requestOptions{URL: "/login"})
	r2 := request(g, requestOptions{
		Method: "POST",
		URL:    "/login",
		Headers: map[string]string{
			"Cookie":       r1.Header().Get("Set-Cookie"),
			"Content-Type": "application/x-www-form-urlencoded",
		},
		Body: strings.NewReader("wtf=" + token),
	})

	if body := r2.Body.String(); body != "OK" {
		t.Error("Response is not OK: ", body)
	}
}
<|endoftext|>


<commit_before>package data

var DEBUG_SIZE int = 100

type Reader struct {
	data []byte
	Cursor uint64
	Size uint64
}

// Creates a new Reader for reading the data from the byte array.
func NewReader (da []byte) *Reader {
	r := new(Reader)
	r.data = da
	r.Cursor = 0
	r.Size = uint64(len(da))
	return r
}

// Reads the number of bytes passed in as size from the data byte array in
// the Reader struct. It then concatinates the bytes and returns them as a
// unsigned integer.
func (r *Reader) Read(size uint) uint {

	var (
		value uint = 0
		i uint64 = 0
	)
	for ; i <uint64(size) ;i++ {
		value |= (uint(r.data[r.Cursor+i]) << ((uint64(size)-i-1)*8))
	}
	r.Cursor += uint64(size)
	return value
}

// Reads the number of bytes passed in as size from the data byte array in the Reader
// struct. It returns a byte array from the cursors current position to the cursor
// plus the size.

func (r *Reader) ReadBytes(size uint64) []byte {
	value:=r.data[r.Cursor:r.Cursor+size]
	r.Cursor += size
	return value
}
<commit_msg>fixed uninitialized and overflow bugs<commit_after>package data

import (
	"log"
)

var DEBUG_SIZE int = 100

type Reader struct {
	data []byte
	Cursor uint64
	Size uint64
}

// Creates a new Reader for reading the data from the byte array.
func NewReader (da []byte) *Reader {
	r := new(Reader)
	r.data = da
	r.Cursor = 0
	r.Size = uint64(len(da))
	return r
}

// Reads the number of bytes passed in as size from the data byte array in
// the Reader struct. It then concatinates the bytes and returns them as a
// unsigned integer.
func (r *Reader) Read(size uint) uint {
  if r.data == nil {
    log.Printf( "attempted to read from null buffer in data.Read()\n" )
    return 0
  }

  var idx uint = uint( r.Cursor ) + size
  var idx_max uint = uint( len( r.data ) )
  var idx_ uint
  if idx <= idx_max {
    idx_ = idx
  } else {
    idx_ = idx_max
  }
  idx_ -= uint( r.Cursor )

  var value uint = 0
  var i uint = 0

  for ; i < idx_ ; i++ {
    value |= uint(uint(r.data[r.Cursor+uint64(i)]) << ((idx_-i-1)*8))
  }
  r.Cursor += uint64( idx_ )
  return value
}

// Reads the number of bytes passed in as size from the data byte array in the Reader
// struct. It returns a byte array from the cursors current position to the cursor
// plus the size.

func (r *Reader) ReadBytes(size uint64) []byte {
	if r.data == nil {
		log.Printf( "attempted to read from null buffer in data.Read()\n" )
		return nil
	}

  var idx uint = uint( r.Cursor ) + uint( size )
	var idx_max uint = uint( len( r.data ) )
	var idx_ uint
	if idx <= idx_max {
		idx_ = idx
	} else {
		idx_ = idx_max
	}
	var n_bytes_read int = int( idx_ - uint( r.Cursor ) )

	value:=r.data[r.Cursor:idx_]
	r.Cursor += uint64( n_bytes_read )
	return value
}
<|endoftext|>


<commit_before>// Package storage provides clients for Microsoft Azure Storage Services.
package storage

import (
	"bytes"
	"encoding/base64"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

const (
	// DefaultBaseURL is the domain name used for storage requests when a
	// default client is created.
	DefaultBaseURL = "core.windows.net"

	// DefaultAPIVersion is the  Azure Storage API version string used when a
	// basic client is created.
	DefaultAPIVersion = "2015-02-21"

	defaultUseHTTPS = true

	blobServiceName  = "blob"
	tableServiceName = "table"
	queueServiceName = "queue"
	fileServiceName  = "file"
)

// Client is the object that needs to be constructed to perform
// operations on the storage account.
type Client struct {
	accountName string
	accountKey  []byte
	useHTTPS    bool
	baseURL     string
	apiVersion  string
}

type storageResponse struct {
	statusCode int
	headers    http.Header
	body       io.ReadCloser
}

// AzureStorageServiceError contains fields of the error response from
// Azure Storage Service REST API. See https://msdn.microsoft.com/en-us/library/azure/dd179382.aspx
// Some fields might be specific to certain calls.
type AzureStorageServiceError struct {
	Code                      string `xml:"Code"`
	Message                   string `xml:"Message"`
	AuthenticationErrorDetail string `xml:"AuthenticationErrorDetail"`
	QueryParameterName        string `xml:"QueryParameterName"`
	QueryParameterValue       string `xml:"QueryParameterValue"`
	Reason                    string `xml:"Reason"`
	StatusCode                int
	RequestID                 string
}

// UnexpectedStatusCodeError is returned when a storage service responds with neither an error
// nor with an HTTP status code indicating success.
type UnexpectedStatusCodeError struct {
	allowed []int
	got     int
}

func (e UnexpectedStatusCodeError) Error() string {
	s := func(i int) string { return fmt.Sprintf("%d %s", i, http.StatusText(i)) }

	got := s(e.got)
	expected := []string{}
	for _, v := range e.allowed {
		expected = append(expected, s(v))
	}
	return fmt.Sprintf("storage: status code from service response is %s; was expecting %s", got, strings.Join(expected, " or "))
}

// Got is the actual status code returned by Azure.
func (e UnexpectedStatusCodeError) Got() int {
	return e.got
}

// NewBasicClient constructs a Client with given storage service name and
// key.
func NewBasicClient(accountName, accountKey string) (Client, error) {
	return NewClient(accountName, accountKey, DefaultBaseURL, DefaultAPIVersion, defaultUseHTTPS)
}

// NewClient constructs a Client. This should be used if the caller wants
// to specify whether to use HTTPS, a specific REST API version or a custom
// storage endpoint than Azure Public Cloud.
func NewClient(accountName, accountKey, blobServiceBaseURL, apiVersion string, useHTTPS bool) (Client, error) {
	var c Client
	if accountName == "" {
		return c, fmt.Errorf("azure: account name required")
	} else if accountKey == "" {
		return c, fmt.Errorf("azure: account key required")
	} else if blobServiceBaseURL == "" {
		return c, fmt.Errorf("azure: base storage service url required")
	}

	key, err := base64.StdEncoding.DecodeString(accountKey)
	if err != nil {
		return c, err
	}

	return Client{
		accountName: accountName,
		accountKey:  key,
		useHTTPS:    useHTTPS,
		baseURL:     blobServiceBaseURL,
		apiVersion:  apiVersion,
	}, nil
}

func (c Client) getBaseURL(service string) string {
	scheme := "http"
	if c.useHTTPS {
		scheme = "https"
	}

	host := fmt.Sprintf("%s.%s.%s", c.accountName, service, c.baseURL)

	u := &url.URL{
		Scheme: scheme,
		Host:   host}
	return u.String()
}

func (c Client) getEndpoint(service, path string, params url.Values) string {
	u, err := url.Parse(c.getBaseURL(service))
	if err != nil {
		// really should not be happening
		panic(err)
	}

	if path == "" {
		path = "/" // API doesn't accept path segments not starting with '/'
	}

	u.Path = path
	u.RawQuery = params.Encode()
	return u.String()
}

// GetBlobService returns a BlobStorageClient which can operate on the blob
// service of the storage account.
func (c Client) GetBlobService() BlobStorageClient {
	return BlobStorageClient{c}
}

// GetQueueService returns a QueueServiceClient which can operate on the queue
// service of the storage account.
func (c Client) GetQueueService() QueueServiceClient {
	return QueueServiceClient{c}
}

// GetFileService returns a FileServiceClient which can operate on the file
// service of the storage account.
func (c Client) GetFileService() FileServiceClient {
	return FileServiceClient{c}
}

func (c Client) createAuthorizationHeader(canonicalizedString string) string {
	signature := c.computeHmac256(canonicalizedString)
	return fmt.Sprintf("%s %s:%s", "SharedKey", c.accountName, signature)
}

func (c Client) getAuthorizationHeader(verb, url string, headers map[string]string) (string, error) {
	canonicalizedResource, err := c.buildCanonicalizedResource(url)
	if err != nil {
		return "", err
	}

	canonicalizedString := c.buildCanonicalizedString(verb, headers, canonicalizedResource)
	return c.createAuthorizationHeader(canonicalizedString), nil
}

func (c Client) getStandardHeaders() map[string]string {
	return map[string]string{
		"x-ms-version": c.apiVersion,
		"x-ms-date":    currentTimeRfc1123Formatted(),
	}
}

func (c Client) buildCanonicalizedHeader(headers map[string]string) string {
	cm := make(map[string]string)

	for k, v := range headers {
		headerName := strings.TrimSpace(strings.ToLower(k))
		match, _ := regexp.MatchString("x-ms-", headerName)
		if match {
			cm[headerName] = v
		}
	}

	if len(cm) == 0 {
		return ""
	}

	keys := make([]string, 0, len(cm))
	for key := range cm {
		keys = append(keys, key)
	}

	sort.Strings(keys)

	ch := ""

	for i, key := range keys {
		if i == len(keys)-1 {
			ch += fmt.Sprintf("%s:%s", key, cm[key])
		} else {
			ch += fmt.Sprintf("%s:%s\n", key, cm[key])
		}
	}
	return ch
}

func (c Client) buildCanonicalizedResource(uri string) (string, error) {
	errMsg := "buildCanonicalizedResource error: %s"
	u, err := url.Parse(uri)
	if err != nil {
		return "", fmt.Errorf(errMsg, err.Error())
	}

	cr := "/" + c.accountName
	if len(u.Path) > 0 {
		cr += u.Path
	}

	params, err := url.ParseQuery(u.RawQuery)
	if err != nil {
		return "", fmt.Errorf(errMsg, err.Error())
	}

	if len(params) > 0 {
		cr += "\n"
		keys := make([]string, 0, len(params))
		for key := range params {
			keys = append(keys, key)
		}

		sort.Strings(keys)

		for i, key := range keys {
			if len(params[key]) > 1 {
				sort.Strings(params[key])
			}

			if i == len(keys)-1 {
				cr += fmt.Sprintf("%s:%s", key, strings.Join(params[key], ","))
			} else {
				cr += fmt.Sprintf("%s:%s\n", key, strings.Join(params[key], ","))
			}
		}
	}
	return cr, nil
}

func (c Client) buildCanonicalizedString(verb string, headers map[string]string, canonicalizedResource string) string {
	contentLength := headers["Content-Length"]
	if contentLength == "0" {
		contentLength = ""
	}
	canonicalizedString := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s",
		verb,
		headers["Content-Encoding"],
		headers["Content-Language"],
		contentLength,
		headers["Content-MD5"],
		headers["Content-Type"],
		headers["Date"],
		headers["If-Modified-Since"],
		headers["If-Match"],
		headers["If-None-Match"],
		headers["If-Unmodified-Since"],
		headers["Range"],
		c.buildCanonicalizedHeader(headers),
		canonicalizedResource)

	return canonicalizedString
}

func (c Client) exec(verb, url string, headers map[string]string, body io.Reader) (*storageResponse, error) {
	authHeader, err := c.getAuthorizationHeader(verb, url, headers)
	if err != nil {
		return nil, err
	}
	headers["Authorization"] = authHeader

	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest(verb, url, body)
	if err != nil {
		return nil, errors.New("azure/storage: error creating request: " + err.Error())
	}

	if clstr, ok := headers["Content-Length"]; ok {
		// content length header is being signed, but completely ignored by golang.
		// instead we have to use the ContentLength property on the request struct
		// (see https://golang.org/src/net/http/request.go?s=18140:18370#L536 and
		// https://golang.org/src/net/http/transfer.go?s=1739:2467#L49)
		req.ContentLength, err = strconv.ParseInt(clstr, 10, 64)
		if err != nil {
			return nil, err
		}
	}
	for k, v := range headers {
		req.Header.Add(k, v)
	}
	httpClient := http.Client{}
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	statusCode := resp.StatusCode
	if statusCode >= 400 && statusCode <= 505 {
		var respBody []byte
		respBody, err = readResponseBody(resp)
		if err != nil {
			return nil, err
		}

		if len(respBody) == 0 {
			// no error in response body
			err = fmt.Errorf("storage: service returned without a response body (%s)", resp.Status)
		} else {
			// response contains storage service error object, unmarshal
			storageErr, errIn := serviceErrFromXML(respBody, resp.StatusCode, resp.Header.Get("x-ms-request-id"))
			if err != nil { // error unmarshaling the error response
				err = errIn
			}
			err = storageErr
		}
		return &storageResponse{
			statusCode: resp.StatusCode,
			headers:    resp.Header,
			body:       ioutil.NopCloser(bytes.NewReader(respBody)), /* restore the body */
		}, err
	}

	return &storageResponse{
		statusCode: resp.StatusCode,
		headers:    resp.Header,
		body:       resp.Body}, nil
}

func readResponseBody(resp *http.Response) ([]byte, error) {
	defer resp.Body.Close()
	out, err := ioutil.ReadAll(resp.Body)
	if err == io.EOF {
		err = nil
	}
	return out, err
}

func serviceErrFromXML(body []byte, statusCode int, requestID string) (AzureStorageServiceError, error) {
	var storageErr AzureStorageServiceError
	if err := xml.Unmarshal(body, &storageErr); err != nil {
		return storageErr, err
	}
	storageErr.StatusCode = statusCode
	storageErr.RequestID = requestID
	return storageErr, nil
}

func (e AzureStorageServiceError) Error() string {
	return fmt.Sprintf("storage: service returned error: StatusCode=%d, ErrorCode=%s, ErrorMessage=%s, RequestId=%s, QueryParameterName=%s, QueryParameterValue=%s",
		e.StatusCode, e.Code, e.Message, e.RequestID, e.QueryParameterName, e.QueryParameterValue)
}

// checkRespCode returns UnexpectedStatusError if the given response code is not
// one of the allowed status codes; otherwise nil.
func checkRespCode(respCode int, allowed []int) error {
	for _, v := range allowed {
		if respCode == v {
			return nil
		}
	}
	return UnexpectedStatusCodeError{allowed, respCode}
}
<commit_msg>Reuse http.Client across requests, and allow caller to provide its own.<commit_after>// Package storage provides clients for Microsoft Azure Storage Services.
package storage

import (
	"bytes"
	"encoding/base64"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

const (
	// DefaultBaseURL is the domain name used for storage requests when a
	// default client is created.
	DefaultBaseURL = "core.windows.net"

	// DefaultAPIVersion is the  Azure Storage API version string used when a
	// basic client is created.
	DefaultAPIVersion = "2015-02-21"

	defaultUseHTTPS = true

	blobServiceName  = "blob"
	tableServiceName = "table"
	queueServiceName = "queue"
	fileServiceName  = "file"
)

// Client is the object that needs to be constructed to perform
// operations on the storage account.
type Client struct {
	// HTTPClient is the http.Client used to initiate API
	// requests.  If it is nil, http.DefaultClient is used.
	HTTPClient *http.Client

	accountName string
	accountKey  []byte
	useHTTPS    bool
	baseURL     string
	apiVersion  string
}

type storageResponse struct {
	statusCode int
	headers    http.Header
	body       io.ReadCloser
}

// AzureStorageServiceError contains fields of the error response from
// Azure Storage Service REST API. See https://msdn.microsoft.com/en-us/library/azure/dd179382.aspx
// Some fields might be specific to certain calls.
type AzureStorageServiceError struct {
	Code                      string `xml:"Code"`
	Message                   string `xml:"Message"`
	AuthenticationErrorDetail string `xml:"AuthenticationErrorDetail"`
	QueryParameterName        string `xml:"QueryParameterName"`
	QueryParameterValue       string `xml:"QueryParameterValue"`
	Reason                    string `xml:"Reason"`
	StatusCode                int
	RequestID                 string
}

// UnexpectedStatusCodeError is returned when a storage service responds with neither an error
// nor with an HTTP status code indicating success.
type UnexpectedStatusCodeError struct {
	allowed []int
	got     int
}

func (e UnexpectedStatusCodeError) Error() string {
	s := func(i int) string { return fmt.Sprintf("%d %s", i, http.StatusText(i)) }

	got := s(e.got)
	expected := []string{}
	for _, v := range e.allowed {
		expected = append(expected, s(v))
	}
	return fmt.Sprintf("storage: status code from service response is %s; was expecting %s", got, strings.Join(expected, " or "))
}

// Got is the actual status code returned by Azure.
func (e UnexpectedStatusCodeError) Got() int {
	return e.got
}

// NewBasicClient constructs a Client with given storage service name and
// key.
func NewBasicClient(accountName, accountKey string) (Client, error) {
	return NewClient(accountName, accountKey, DefaultBaseURL, DefaultAPIVersion, defaultUseHTTPS)
}

// NewClient constructs a Client. This should be used if the caller wants
// to specify whether to use HTTPS, a specific REST API version or a custom
// storage endpoint than Azure Public Cloud.
func NewClient(accountName, accountKey, blobServiceBaseURL, apiVersion string, useHTTPS bool) (Client, error) {
	var c Client
	if accountName == "" {
		return c, fmt.Errorf("azure: account name required")
	} else if accountKey == "" {
		return c, fmt.Errorf("azure: account key required")
	} else if blobServiceBaseURL == "" {
		return c, fmt.Errorf("azure: base storage service url required")
	}

	key, err := base64.StdEncoding.DecodeString(accountKey)
	if err != nil {
		return c, err
	}

	return Client{
		accountName: accountName,
		accountKey:  key,
		useHTTPS:    useHTTPS,
		baseURL:     blobServiceBaseURL,
		apiVersion:  apiVersion,
	}, nil
}

func (c Client) getBaseURL(service string) string {
	scheme := "http"
	if c.useHTTPS {
		scheme = "https"
	}

	host := fmt.Sprintf("%s.%s.%s", c.accountName, service, c.baseURL)

	u := &url.URL{
		Scheme: scheme,
		Host:   host}
	return u.String()
}

func (c Client) getEndpoint(service, path string, params url.Values) string {
	u, err := url.Parse(c.getBaseURL(service))
	if err != nil {
		// really should not be happening
		panic(err)
	}

	if path == "" {
		path = "/" // API doesn't accept path segments not starting with '/'
	}

	u.Path = path
	u.RawQuery = params.Encode()
	return u.String()
}

// GetBlobService returns a BlobStorageClient which can operate on the blob
// service of the storage account.
func (c Client) GetBlobService() BlobStorageClient {
	return BlobStorageClient{c}
}

// GetQueueService returns a QueueServiceClient which can operate on the queue
// service of the storage account.
func (c Client) GetQueueService() QueueServiceClient {
	return QueueServiceClient{c}
}

// GetFileService returns a FileServiceClient which can operate on the file
// service of the storage account.
func (c Client) GetFileService() FileServiceClient {
	return FileServiceClient{c}
}

func (c Client) createAuthorizationHeader(canonicalizedString string) string {
	signature := c.computeHmac256(canonicalizedString)
	return fmt.Sprintf("%s %s:%s", "SharedKey", c.accountName, signature)
}

func (c Client) getAuthorizationHeader(verb, url string, headers map[string]string) (string, error) {
	canonicalizedResource, err := c.buildCanonicalizedResource(url)
	if err != nil {
		return "", err
	}

	canonicalizedString := c.buildCanonicalizedString(verb, headers, canonicalizedResource)
	return c.createAuthorizationHeader(canonicalizedString), nil
}

func (c Client) getStandardHeaders() map[string]string {
	return map[string]string{
		"x-ms-version": c.apiVersion,
		"x-ms-date":    currentTimeRfc1123Formatted(),
	}
}

func (c Client) buildCanonicalizedHeader(headers map[string]string) string {
	cm := make(map[string]string)

	for k, v := range headers {
		headerName := strings.TrimSpace(strings.ToLower(k))
		match, _ := regexp.MatchString("x-ms-", headerName)
		if match {
			cm[headerName] = v
		}
	}

	if len(cm) == 0 {
		return ""
	}

	keys := make([]string, 0, len(cm))
	for key := range cm {
		keys = append(keys, key)
	}

	sort.Strings(keys)

	ch := ""

	for i, key := range keys {
		if i == len(keys)-1 {
			ch += fmt.Sprintf("%s:%s", key, cm[key])
		} else {
			ch += fmt.Sprintf("%s:%s\n", key, cm[key])
		}
	}
	return ch
}

func (c Client) buildCanonicalizedResource(uri string) (string, error) {
	errMsg := "buildCanonicalizedResource error: %s"
	u, err := url.Parse(uri)
	if err != nil {
		return "", fmt.Errorf(errMsg, err.Error())
	}

	cr := "/" + c.accountName
	if len(u.Path) > 0 {
		cr += u.Path
	}

	params, err := url.ParseQuery(u.RawQuery)
	if err != nil {
		return "", fmt.Errorf(errMsg, err.Error())
	}

	if len(params) > 0 {
		cr += "\n"
		keys := make([]string, 0, len(params))
		for key := range params {
			keys = append(keys, key)
		}

		sort.Strings(keys)

		for i, key := range keys {
			if len(params[key]) > 1 {
				sort.Strings(params[key])
			}

			if i == len(keys)-1 {
				cr += fmt.Sprintf("%s:%s", key, strings.Join(params[key], ","))
			} else {
				cr += fmt.Sprintf("%s:%s\n", key, strings.Join(params[key], ","))
			}
		}
	}
	return cr, nil
}

func (c Client) buildCanonicalizedString(verb string, headers map[string]string, canonicalizedResource string) string {
	contentLength := headers["Content-Length"]
	if contentLength == "0" {
		contentLength = ""
	}
	canonicalizedString := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s",
		verb,
		headers["Content-Encoding"],
		headers["Content-Language"],
		contentLength,
		headers["Content-MD5"],
		headers["Content-Type"],
		headers["Date"],
		headers["If-Modified-Since"],
		headers["If-Match"],
		headers["If-None-Match"],
		headers["If-Unmodified-Since"],
		headers["Range"],
		c.buildCanonicalizedHeader(headers),
		canonicalizedResource)

	return canonicalizedString
}

func (c Client) exec(verb, url string, headers map[string]string, body io.Reader) (*storageResponse, error) {
	authHeader, err := c.getAuthorizationHeader(verb, url, headers)
	if err != nil {
		return nil, err
	}
	headers["Authorization"] = authHeader

	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest(verb, url, body)
	if err != nil {
		return nil, errors.New("azure/storage: error creating request: " + err.Error())
	}

	if clstr, ok := headers["Content-Length"]; ok {
		// content length header is being signed, but completely ignored by golang.
		// instead we have to use the ContentLength property on the request struct
		// (see https://golang.org/src/net/http/request.go?s=18140:18370#L536 and
		// https://golang.org/src/net/http/transfer.go?s=1739:2467#L49)
		req.ContentLength, err = strconv.ParseInt(clstr, 10, 64)
		if err != nil {
			return nil, err
		}
	}
	for k, v := range headers {
		req.Header.Add(k, v)
	}

	httpClient := c.HTTPClient
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	statusCode := resp.StatusCode
	if statusCode >= 400 && statusCode <= 505 {
		var respBody []byte
		respBody, err = readResponseBody(resp)
		if err != nil {
			return nil, err
		}

		if len(respBody) == 0 {
			// no error in response body
			err = fmt.Errorf("storage: service returned without a response body (%s)", resp.Status)
		} else {
			// response contains storage service error object, unmarshal
			storageErr, errIn := serviceErrFromXML(respBody, resp.StatusCode, resp.Header.Get("x-ms-request-id"))
			if err != nil { // error unmarshaling the error response
				err = errIn
			}
			err = storageErr
		}
		return &storageResponse{
			statusCode: resp.StatusCode,
			headers:    resp.Header,
			body:       ioutil.NopCloser(bytes.NewReader(respBody)), /* restore the body */
		}, err
	}

	return &storageResponse{
		statusCode: resp.StatusCode,
		headers:    resp.Header,
		body:       resp.Body}, nil
}

func readResponseBody(resp *http.Response) ([]byte, error) {
	defer resp.Body.Close()
	out, err := ioutil.ReadAll(resp.Body)
	if err == io.EOF {
		err = nil
	}
	return out, err
}

func serviceErrFromXML(body []byte, statusCode int, requestID string) (AzureStorageServiceError, error) {
	var storageErr AzureStorageServiceError
	if err := xml.Unmarshal(body, &storageErr); err != nil {
		return storageErr, err
	}
	storageErr.StatusCode = statusCode
	storageErr.RequestID = requestID
	return storageErr, nil
}

func (e AzureStorageServiceError) Error() string {
	return fmt.Sprintf("storage: service returned error: StatusCode=%d, ErrorCode=%s, ErrorMessage=%s, RequestId=%s, QueryParameterName=%s, QueryParameterValue=%s",
		e.StatusCode, e.Code, e.Message, e.RequestID, e.QueryParameterName, e.QueryParameterValue)
}

// checkRespCode returns UnexpectedStatusError if the given response code is not
// one of the allowed status codes; otherwise nil.
func checkRespCode(respCode int, allowed []int) error {
	for _, v := range allowed {
		if respCode == v {
			return nil
		}
	}
	return UnexpectedStatusCodeError{allowed, respCode}
}
<|endoftext|>


<commit_before>package storage

import (
	"io"

	raw "code.google.com/p/google-api-go-client/storage/v1beta2"
)

// contentTyper implements ContentTyper to enable an
// io.ReadCloser to specify its MIME type.
type contentTyper struct {
	io.ReadCloser
	t string
}

func (c *contentTyper) ContentType() string {
	return c.t
}

// newObjectWriter returns a new objectWriter that writes to
// the file that is specified by info.Bucket and info.Name.
// Metadata changes are also reflected on the remote object
// entity, read-only fields are ignored during the write operation.
func newObjectWriter(conn *conn, info *ObjectInfo) *objectWriter {
	w := &objectWriter{
		conn: conn,
		info: info,
	}
	pr, pw := io.Pipe()
	w.rc = &contentTyper{pr, info.ContentType}
	w.pw = pw
	go func() {
		// TODO(jbd): Return the inserted/updated object entity.
		_, w.err = conn.s.Objects.Insert(
			info.Bucket, info.toRawObject()).Media(w.rc).Do()
	}()
	return w
}

// objectWriter is an io.WriteCloser that opens a connection
// to update the metadata and file contents of a GCS object.
type objectWriter struct {
	conn *conn
	info *ObjectInfo

	rc  io.ReadCloser
	pw  *io.PipeWriter
	err error
}

// Write writes len(p) bytes to the object. It returns the number
// of the bytes written, or an error if there is a problem occured
// during the write. It's a blocking operation, and will not return
// until the bytes are written to the underlying socket.
func (w *objectWriter) Write(p []byte) (n int, err error) {
	if w.err != nil {
		return 0, w.err
	}
	return w.pw.Write(p)
}

// Close closes the writer and cleans up other resources
// used by the writer.
func (w *objectWriter) Close() error {
	if w.err != nil {
		return w.err
	}
	w.rc.Close()
	return w.pw.Close()
}

// ObjectInfo represents a Google Cloud Storage (GCS) object.
type ObjectInfo struct {
	// Bucket is the name of the bucket containing this GCS object.
	Bucket string `json:"bucket,omitempty"`

	// Name is the name of the object.
	Name string `json:"name,omitempty"`

	// ContentType is the MIME type of the object's content.
	ContentType string `json:"contentType,omitempty"`

	// Size is the length of the object's content.
	// Read-only.
	Size uint64 `json:"size,omitempty"`

	// ContentEncoding is the encoding of the object's content.
	// Read-only.
	ContentEncoding string `json:"contentEncoding,omitempty"`

	// MD5 is the MD5 hash of the data.
	// Read-only.
	MD5 []byte `json:"md5Hash,omitempty"`

	// CRC32C is the CRC32C checksum of the object's content.
	// Read-only.
	CRC32C []byte `json:"crc32c,omitempty"`

	// MediaLink is an URL to the object's content.
	// Read-only.
	MediaLink string `json:"mediaLink,omitempty"`

	// Metadata represents user-provided metadata, in key/value pairs.
	// It can be nil if no metadata is provided.
	Metadata map[string]string `json:"metadata,omitempty"`

	// Generation is the generation version of the object's content.
	// Read-only.
	Generation int64 `json:"generation,omitempty"`

	// MetaGeneration is the version of the metadata for this
	// object at this generation. This field is used for preconditions
	// and for detecting changes in metadata. A metageneration number
	// is only meaningful in the context of a particular generation
	// of a particular object.
	// Read-only.
	MetaGeneration int64 `json:"metageneration,omitempty"`

	// TODO(jbd): Add ACL and owner.
	// TODO(jbd): Add timeDelete and updated.
}

func (o *ObjectInfo) toRawObject() *raw.Object {
	// TODO(jbd): add ACL and owner
	return &raw.Object{
		Bucket: o.Bucket,
		Name:   o.Name,
	}
}

func newObjectInfo(o *raw.Object) *ObjectInfo {
	if o == nil {
		return nil
	}
	return &ObjectInfo{
		Bucket:          o.Bucket,
		Name:            o.Name,
		ContentType:     o.ContentType,
		ContentEncoding: o.ContentEncoding,
		Size:            o.Size,
		MD5:             []byte(o.Md5Hash),
		CRC32C:          []byte(o.Crc32c),
		MediaLink:       o.MediaLink,
		Generation:      o.Generation,
		MetaGeneration:  o.Metageneration,
	}
}
<commit_msg>storage: Adding access rule controls and owner info.<commit_after>package storage

import (
	"io"

	raw "code.google.com/p/google-api-go-client/storage/v1beta2"
)

// contentTyper implements ContentTyper to enable an
// io.ReadCloser to specify its MIME type.
type contentTyper struct {
	io.ReadCloser
	t string
}

func (c *contentTyper) ContentType() string {
	return c.t
}

// newObjectWriter returns a new objectWriter that writes to
// the file that is specified by info.Bucket and info.Name.
// Metadata changes are also reflected on the remote object
// entity, read-only fields are ignored during the write operation.
func newObjectWriter(conn *conn, info *ObjectInfo) *objectWriter {
	w := &objectWriter{
		conn: conn,
		info: info,
	}
	pr, pw := io.Pipe()
	w.rc = &contentTyper{pr, info.ContentType}
	w.pw = pw
	go func() {
		// TODO(jbd): Return the inserted/updated object entity.
		_, w.err = conn.s.Objects.Insert(
			info.Bucket, info.toRawObject()).Media(w.rc).Do()
	}()
	return w
}

// objectWriter is an io.WriteCloser that opens a connection
// to update the metadata and file contents of a GCS object.
type objectWriter struct {
	conn *conn
	info *ObjectInfo

	rc  io.ReadCloser
	pw  *io.PipeWriter
	err error
}

// Write writes len(p) bytes to the object. It returns the number
// of the bytes written, or an error if there is a problem occured
// during the write. It's a blocking operation, and will not return
// until the bytes are written to the underlying socket.
func (w *objectWriter) Write(p []byte) (n int, err error) {
	if w.err != nil {
		return 0, w.err
	}
	return w.pw.Write(p)
}

// Close closes the writer and cleans up other resources
// used by the writer.
func (w *objectWriter) Close() error {
	if w.err != nil {
		return w.err
	}
	w.rc.Close()
	return w.pw.Close()
}

// AccessControlRule represents an access control rule entry for
// a Google Cloud Storage object.
type AccessControlRule struct {
	// Entity identifies the entity holding the current
	// rule's permissions. It could be in the form of:
	// - "user-<userId>"
	// - "user-<email>"
	// - "group-<groupId>"
	// - "group-<email>"
	// - "domain-<domain>"
	// - "project-team-<projectId>"
	// - "allUsers"
	// - "allAuthenticatedUsers"
	Entity string `json:"entity,omitempty"`

	// Role is the the access permission for the entity. Can be READER or OWNER.
	Role string `json:"role,omitempty"`
}

// Owner represents the owner of a GCS object.
type Owner struct {
	// Entity identifies the owner, it's always in the form of "user-<userId>".
	Entity string `json:"entity,omitempty"`
}

// ObjectInfo represents a Google Cloud Storage (GCS) object.
type ObjectInfo struct {
	// Bucket is the name of the bucket containing this GCS object.
	Bucket string `json:"bucket,omitempty"`

	// Name is the name of the object.
	Name string `json:"name,omitempty"`

	// ContentType is the MIME type of the object's content.
	ContentType string `json:"contentType,omitempty"`

	// ACL is the access control rule list for the object.
	ACL []*AccessControlRule `json:"acl,omitempty"`

	// Owner is the owner of the object. Owner is alway the original
	// uploader of the object.
	// Read-only.
	Owner *Owner `json:"owner,omitempty"`

	// Size is the length of the object's content.
	// Read-only.
	Size uint64 `json:"size,omitempty"`

	// ContentEncoding is the encoding of the object's content.
	// Read-only.
	ContentEncoding string `json:"contentEncoding,omitempty"`

	// MD5 is the MD5 hash of the data.
	// Read-only.
	MD5 []byte `json:"md5Hash,omitempty"`

	// CRC32C is the CRC32C checksum of the object's content.
	// Read-only.
	CRC32C []byte `json:"crc32c,omitempty"`

	// MediaLink is an URL to the object's content.
	// Read-only.
	MediaLink string `json:"mediaLink,omitempty"`

	// Metadata represents user-provided metadata, in key/value pairs.
	// It can be nil if no metadata is provided.
	Metadata map[string]string `json:"metadata,omitempty"`

	// Generation is the generation version of the object's content.
	// Read-only.
	Generation int64 `json:"generation,omitempty"`

	// MetaGeneration is the version of the metadata for this
	// object at this generation. This field is used for preconditions
	// and for detecting changes in metadata. A metageneration number
	// is only meaningful in the context of a particular generation
	// of a particular object.
	// Read-only.
	MetaGeneration int64 `json:"metageneration,omitempty"`

	// TODO(jbd): Add timeDelete and updated.
}

func (o *ObjectInfo) toRawObject() *raw.Object {
	acl := make([]*raw.ObjectAccessControl, len(o.ACL))
	for i, rule := range o.ACL {
		acl[i] = &raw.ObjectAccessControl{
			Entity: rule.Entity,
			Role:   rule.Role,
		}
	}
	return &raw.Object{
		Bucket: o.Bucket,
		Name:   o.Name,
		Acl:    acl,
	}
}

func newObjectInfo(o *raw.Object) *ObjectInfo {
	if o == nil {
		return nil
	}
	acl := make([]*AccessControlRule, len(o.Acl))
	for i, rule := range o.Acl {
		acl[i] = &AccessControlRule{
			Entity: rule.Entity,
			Role:   rule.Role,
		}
	}
	return &ObjectInfo{
		Bucket:          o.Bucket,
		Name:            o.Name,
		ContentType:     o.ContentType,
		ACL:             acl,
		Owner:           &Owner{Entity: o.Owner.Entity},
		ContentEncoding: o.ContentEncoding,
		Size:            o.Size,
		MD5:             []byte(o.Md5Hash),
		CRC32C:          []byte(o.Crc32c),
		MediaLink:       o.MediaLink,
		Generation:      o.Generation,
		MetaGeneration:  o.Metageneration,
	}
}
<|endoftext|>


<commit_before>package store

import (
	log "github.com/Sirupsen/logrus"

	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"
)

const postgresGormLogMode = false

type PostgresKVStore struct {
	*kvStore
	config PostgresConfig
}

func NewPostgresKVStore(postgresConfig PostgresConfig) *PostgresKVStore {
	return &PostgresKVStore{
		kvStore: &kvStore{logger: log.WithFields(log.Fields{"module": "kv-postgres"})},
		config:  postgresConfig,
	}
}

func (kvStore *PostgresKVStore) Open() error {
	logger := kvStore.logger.WithField("config", kvStore.config)
	logger.Info("Opening database")

	gormdb, err := gorm.Open("postgres", kvStore.config.connectionString())
	if err != nil {
		logger.WithField("err", err).Error("Error opening database")
		return err
	}

	if err := gormdb.DB().Ping(); err != nil {
		kvStore.logger.WithError(err).Error("Error pinging database")
	} else {
		kvStore.logger.Info("Ping reply from database")
	}

	gormdb.LogMode(postgresGormLogMode)
	gormdb.SingularTable(true)
	gormdb.DB().SetMaxIdleConns(kvStore.config.MaxIdleConns)
	gormdb.DB().SetMaxOpenConns(kvStore.config.MaxOpenConns)
	if err := gormdb.AutoMigrate(&kvEntry{}).Error; err != nil {
		logger.WithField("err", err).Error("Error in schema migration")
		return err
	}

	logger.Info("Ensured database schema")
	kvStore.db = gormdb
	return nil
}
<commit_msg>adding doc<commit_after>package store

import (
	log "github.com/Sirupsen/logrus"

	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"
)

const postgresGormLogMode = false

// PostgresKVStore extends a gorm-based kvStore with a Postgresql-specific configuration.
type PostgresKVStore struct {
	*kvStore
	config PostgresConfig
}

// NewPostgresKVStore returns a new configured PostgresKVStore (not opened yet).
func NewPostgresKVStore(postgresConfig PostgresConfig) *PostgresKVStore {
	return &PostgresKVStore{
		kvStore: &kvStore{logger: log.WithFields(log.Fields{"module": "kv-postgres"})},
		config:  postgresConfig,
	}
}

// Open a connection to Postgresql database, or return an error.
func (kvStore *PostgresKVStore) Open() error {
	logger := kvStore.logger.WithField("config", kvStore.config)
	logger.Info("Opening database")

	gormdb, err := gorm.Open("postgres", kvStore.config.connectionString())
	if err != nil {
		logger.WithField("err", err).Error("Error opening database")
		return err
	}

	if err := gormdb.DB().Ping(); err != nil {
		kvStore.logger.WithError(err).Error("Error pinging database")
	} else {
		kvStore.logger.Info("Ping reply from database")
	}

	gormdb.LogMode(postgresGormLogMode)
	gormdb.SingularTable(true)
	gormdb.DB().SetMaxIdleConns(kvStore.config.MaxIdleConns)
	gormdb.DB().SetMaxOpenConns(kvStore.config.MaxOpenConns)
	if err := gormdb.AutoMigrate(&kvEntry{}).Error; err != nil {
		logger.WithField("err", err).Error("Error in schema migration")
		return err
	}

	logger.Info("Ensured database schema")
	kvStore.db = gormdb
	return nil
}
<|endoftext|>


<commit_before>/* Copyright (c) 2014-2015, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"encoding/json"
	"flag"
	"html/template"
	"log"
	"net/http"
	"strings"
)

var (
	listen = flag.String("l", ":8080", "Host and port to listen to")
)

type Typed struct {
	Type string `json:"type"`
}

type handler struct {
	tmpl   *template.Template
	places []Place
	get    endpoints
}

type endpoints map[string]func(http.ResponseWriter, *http.Request)

func newHttpHandler() *handler {
	h := &handler{
		tmpl: template.Must(template.ParseFiles("index.html")),
	}
	h.places = append(h.places, getClubs()...)
	h.get = endpoints{
		"/":              h.index,
		"/places":        h.getplaces,
		"/places-pretty": h.getplacespretty,
	}
	return h
}

func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	p := strings.ToLower(r.URL.Path)
	switch r.Method {
	case "GET":
		f, e := h.get[p]
		if !e {
			http.Error(w, "unknown endpoint", http.StatusBadRequest)
			return
		}
		f(w, r)
	default:
		http.Error(w, "unsupported action", http.StatusBadRequest)
	}
}

func (h *handler) index(w http.ResponseWriter, r *http.Request) {
	if err := h.tmpl.ExecuteTemplate(w, "index.html", struct {
		Get endpoints
	}{
		Get: h.get,
	}); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (h *handler) getplaces(w http.ResponseWriter, r *http.Request) {
	e := json.NewEncoder(w)
	e.Encode(&h.places)
}

func (h *handler) getplacespretty(w http.ResponseWriter, r *http.Request) {
	b, err := json.MarshalIndent(&h.places, "", "  ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.Write(b)
}

func main() {
	flag.Parse()
	handler := newHttpHandler()
	log.Printf("listen = %s", *listen)
	http.Handle("/", handler)
	log.Println("Up and running!")
	log.Fatal(http.ListenAndServe(*listen, nil))
}
<commit_msg>Weird header<commit_after>/* Copyright (c) 2014-2015, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"encoding/json"
	"flag"
	"html/template"
	"log"
	"net/http"
	"strings"
)

var (
	listen = flag.String("l", ":8080", "Host and port to listen to")
)

type Typed struct {
	Type string `json:"type"`
}

type handler struct {
	tmpl   *template.Template
	places []Place
	get    endpoints
}

type endpoints map[string]func(http.ResponseWriter, *http.Request)

func newHttpHandler() *handler {
	h := &handler{
		tmpl: template.Must(template.ParseFiles("index.html")),
	}
	h.places = append(h.places, getClubs()...)
	h.get = endpoints{
		"/":              h.index,
		"/places":        h.getplaces,
		"/places-pretty": h.getplacespretty,
	}
	return h
}

func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	p := strings.ToLower(r.URL.Path)
	switch r.Method {
	case "GET":
		f, e := h.get[p]
		if !e {
			http.Error(w, "unknown endpoint", http.StatusBadRequest)
			return
		}
		w.Header().Set("Access-Control-Allow-Origin", "*")
		f(w, r)
	default:
		http.Error(w, "unsupported action", http.StatusBadRequest)
	}
}

func (h *handler) index(w http.ResponseWriter, r *http.Request) {
	if err := h.tmpl.ExecuteTemplate(w, "index.html", struct {
		Get endpoints
	}{
		Get: h.get,
	}); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (h *handler) getplaces(w http.ResponseWriter, r *http.Request) {
	e := json.NewEncoder(w)
	e.Encode(&h.places)
}

func (h *handler) getplacespretty(w http.ResponseWriter, r *http.Request) {
	b, err := json.MarshalIndent(&h.places, "", "  ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.Write(b)
}

func main() {
	flag.Parse()
	handler := newHttpHandler()
	log.Printf("listen = %s", *listen)
	http.Handle("/", handler)
	log.Println("Up and running!")
	log.Fatal(http.ListenAndServe(*listen, nil))
}
<|endoftext|>


<commit_before>package ratelimit

import (
	"strconv"
	"testing"
	"time"

	"golang.org/x/time/rate"
)

func TestAllowNAndRetryAfter(t *testing.T) {

	tests := []struct {
		keys             []string
		rule             Rule
		qps              int
		shouldRetryAfter time.Duration
		shouldErr        bool
		expected         bool
	}{
		{
			[]string{"127.0.0.1", "get", "/"}, Rule{Methods: "get", Rate: 2, Burst: 2, Unit: "second"}, 2, 0, false, true,
		},
		{
			[]string{"127.0.0.1", "get,post", "/"}, Rule{Methods: "get,post", Rate: 1, Burst: 2, Unit: "minute"}, 2, 30 * time.Second, false, true,
		},
		{
			[]string{"127.0.0.1", "*", "/"}, Rule{Methods: "*", Rate: 1, Burst: 0, Unit: "hour"}, 1, rate.InfDuration, false, false,
		},
		{
			[]string{"127.0.0.1", "", "/"}, Rule{Methods: "", Rate: 0, Burst: 0}, 2, 0, false, true,
		},
	}

	for i, test := range tests {
		test.keys = append(test.keys, strconv.Itoa(i))
		actual := cl.AllowN(test.keys, test.rule, test.qps)
		retryAfter := cl.RetryAfter(test.keys)
		if retryAfter < test.shouldRetryAfter {
			t.Errorf("E! test %d: shouldRetryAfter %d, got %d", i, test.shouldRetryAfter, retryAfter)
		}
		if actual != test.expected {
			t.Errorf("E! test %d: expected %t, got %t", i, test.expected, actual)
		}
	}

	// spawn multiple goroutines to test concurrent read/write in map
	num := make([]int, 1000)
	for range num {
		go func() {
			for {
				cl.AllowN(tests[0].keys, tests[0].rule, tests[0].qps)
			}
		}()
		go func() {
			for {
				cl.AllowN(tests[0].keys, tests[0].rule, tests[0].qps)
			}
		}()
	}
}

func BenchmarkSingleKey(b *testing.B) {

	keys := []string{"127.0.0.1", "get", "/"}
	for n := 0; n < b.N; n++ {
		benchmarkAllowNAndRetryAfter(keys)
	}
}

func BenchmarkRandomKey(b *testing.B) {

	for n := 0; n < b.N; n++ {
		keys := []string{"127.0.0.1", "get", "/" + strconv.Itoa(n)}
		benchmarkAllowNAndRetryAfter(keys)
	}
}

func benchmarkAllowNAndRetryAfter(keys []string) {

	cl.AllowN(keys, Rule{Methods: "get", Rate: 2, Burst: 2, Unit: "second"}, 1)
	cl.RetryAfter(keys)
}
<commit_msg>Run multiple benchmark<commit_after>package ratelimit

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	"golang.org/x/time/rate"
)

func TestAllowNAndRetryAfter(t *testing.T) {

	tests := []struct {
		keys             []string
		rule             Rule
		qps              int
		shouldRetryAfter time.Duration
		shouldErr        bool
		expected         bool
	}{
		{
			[]string{"127.0.0.1", "get", "/"}, Rule{Methods: "get", Rate: 2, Burst: 2, Unit: "second"}, 2, 0, false, true,
		},
		{
			[]string{"127.0.0.1", "get,post", "/"}, Rule{Methods: "get,post", Rate: 1, Burst: 2, Unit: "minute"}, 2, 30 * time.Second, false, true,
		},
		{
			[]string{"127.0.0.1", "*", "/"}, Rule{Methods: "*", Rate: 1, Burst: 0, Unit: "hour"}, 1, rate.InfDuration, false, false,
		},
		{
			[]string{"127.0.0.1", "", "/"}, Rule{Methods: "", Rate: 0, Burst: 0}, 2, 0, false, true,
		},
	}

	for i, test := range tests {
		test.keys = append(test.keys, strconv.Itoa(i))
		actual := cl.AllowN(test.keys, test.rule, test.qps)
		retryAfter := cl.RetryAfter(test.keys)
		if retryAfter < test.shouldRetryAfter {
			t.Errorf("E! test %d: shouldRetryAfter %d, got %d", i, test.shouldRetryAfter, retryAfter)
		}
		if actual != test.expected {
			t.Errorf("E! test %d: expected %t, got %t", i, test.expected, actual)
		}
	}

	// spawn multiple goroutines to test concurrent read/write in map
	num := make([]int, 1000)
	for range num {
		go func() {
			for {
				cl.AllowN(tests[0].keys, tests[0].rule, tests[0].qps)
			}
		}()
		go func() {
			for {
				cl.AllowN(tests[0].keys, tests[0].rule, tests[0].qps)
			}
		}()
	}
}

func BenchmarkSingleKey(b *testing.B) {

	keys := []string{"127.0.0.1", "get", "/"}
	for i := 1; i <= 8; i *= 2 {
		b.Run(fmt.Sprintf("%d", i), func(b *testing.B) {
			for n := 0; n < b.N; n++ {
				benchmarkAllowNAndRetryAfter(keys)
			}
		})
	}
}

func BenchmarkRandomKey(b *testing.B) {

	for i := 1; i <= 8; i *= 2 {
		b.Run(fmt.Sprintf("%d", i), func(b *testing.B) {
			for n := 0; n < b.N; n++ {
				keys := []string{"127.0.0.1", "get", "/" + strconv.Itoa(i) + "-" + strconv.Itoa(n)}
				benchmarkAllowNAndRetryAfter(keys)
			}
		})
	}
}

func benchmarkAllowNAndRetryAfter(keys []string) {

	cl.AllowN(keys, Rule{Methods: "get", Rate: 2, Burst: 2, Unit: "second"}, 1)
	cl.RetryAfter(keys)
}
<|endoftext|>


<commit_before>// Package calendar generates and has utlilty functions to generate an HTML calendar for use in nlgids.
package calendar

import (
	"fmt"
	"sort"
	"time"
)

// TODO: next button should have a ref to point to something.

var avail = [...]string{"past", "busy", "free"}

type Available int

const (
	past Available = iota
	busy
	free
)

func (a Available) String() string { return avail[a] }

// Calendar holds the HTML that makes up the calendar. Each
// day is indexed by the 12 o' clock night time time.Time.
// All date are in the UTC timezone.
type Calendar struct {
	days  map[time.Time]Available
	begin time.Time
	end   time.Time
	// month we're in (mostly)?
}

type times []time.Time

func (t times) Len() int           { return len(t) }
func (t times) Less(i, j int) bool { return t[i].Before(t[j]) }
func (t times) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }

func (c *Calendar) heading() string {
	// lang!
	s := `

`
	s += "
Sun	Mon	Tue	Wed	Thu	Fri	Sat
\n"
	return s
}

func (c *Calendar) Header() string {
	// Template on Calendar that gets the month from c (or we set it).
	s := `

bla

`
	return s
}

func (c *Calendar) Footer() string {
	return `

`
}

func (c *Calendar) openTR() string  { return "
\n" }
func (c *Calendar) closeTR() string { return "
\n" }

func (c *Calendar) entry(t time.Time) string {
	d := c.days[t]
	day := fmt.Sprintf("%02d", t.Day())
	class := fmt.Sprintf("\t", d)
	close := "	\n"
	href := ""
	switch d {
	case free:
		date := fmt.Sprintf("%4d-%02d-%02d", t.Year(), t.Month(), t.Day())
		href = fmt.Sprintf("%d", date, t.Day()) // SetDate is defined on the page/form itself
		class = fmt.Sprintf("\t", d)
	case busy:
		href = day
	case past:
		href = day
	}
	s := class + href + close
	return s
}

func (c *Calendar) HTML() string {
	s := c.Header()
	s += c.html()
	s += c.Footer()
	return s
}

func (c *Calendar) sort() times {
	keys := times{}
	for k := range c.days {
		keys = append(keys, k)
	}

	sort.Sort(keys)
	return keys
}

func (c *Calendar) html() string {
	keys := c.sort()

	s := c.heading()
	i := 0
	for _, k := range keys {
		if i%7 == 0 {
			if i > 0 {
				s += c.closeTR()
			}
			s += c.openTR()
		}
		s += c.entry(k)
		i++
	}
	s += c.closeTR()
	return s
}

// New creates a new month calendar based on d, d must be in the form: YYYY-MM-DD.
// D can also be the empty string, then the current date is assumed.
func New(d string) (*Calendar, error) {
	date, now := time.Now(), time.Now()
	if d != "" {
		var err error
		date, err = time.Parse("2006-01-02", d)
		if err != nil {
			return nil, err
		}
	}

	cal := &Calendar{days: make(map[time.Time]Available)}

	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
	first := time.Date(date.Year(), date.Month(), 1, 0, 0, 0, 0, time.UTC)
	last := time.Date(date.Year(), date.Month()+1, 1, 0, 0, 0, 0, time.UTC)
	last = last.Add(-24 * time.Hour)

	// Add the remaining days of the previous month.
	for i := 0; i < int(first.Weekday()); i++ {
		lastMonthDay := first.AddDate(0, 0, -1*(i+1))
		cal.days[lastMonthDay] = free

		if lastMonthDay.Before(today) {
			cal.days[lastMonthDay] = past
		}
	}

	// Loop from i to lastDay and add the entire month.
	for i := 1; i <= last.Day(); i++ {
		day := time.Date(date.Year(), date.Month(), i, 0, 0, 0, 0, time.UTC)

		cal.days[day] = free

		if day.Before(today) {
			cal.days[day] = past
		}
	}

	// These are dates in the new month.
	j := 1
	for i := int(last.Weekday()) + 1; i < 7; i++ {
		nextMonthDay := last.AddDate(0, 0, j)
		cal.days[nextMonthDay] = free

		if nextMonthDay.Before(today) {
			cal.days[nextMonthDay] = past
		}

		j++
	}
	times := cal.sort()
	if len(times) > 0 {
		cal.begin = times[0]
		cal.end = times[len(times)-1]
	}

	return cal, nil
}
<commit_msg>more calendar<commit_after>// Package calendar generates and has utlilty functions to generate an HTML calendar for use in nlgids.
package calendar

import (
	"fmt"
	"sort"
	"time"
)

var (
	avail   = [...]string{"past", "busy", "free"}
	monthNL = [...]string{"januari", "februari", "maart", "april", "mei", "juni", "juli", "augustus", "september", "oktober", "november", "december"}
)

type Available int

const (
	past Available = iota
	busy
	free
)

func (a Available) String() string { return avail[a] }

// Calendar holds the HTML that makes up the calendar. Each
// day is indexed by the 12 o' clock night time time.Time.
// All date are in the UTC timezone.
type Calendar struct {
	days  map[time.Time]Available
	begin time.Time
	end   time.Time
	start time.Time // generated for this date
}

type times []time.Time

func (t times) Len() int           { return len(t) }
func (t times) Less(i, j int) bool { return t[i].Before(t[j]) }
func (t times) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }

func (c *Calendar) heading() string {
	s := `

`
	s += "
Sun	Mon	Tue	Wed	Thu	Fri	Sat
\n"
	s += "
zon	maa	din	woe	don	vrij	zat
\n"
	return s
}

// Header returns the header of the calendar.
func (c *Calendar) Header() string {
	month := c.start.Month()

	s := `

`

	s += "
" + month.String() + "
"
	s += "
" + monthNL[month] + "
"

	s += `

`
	return s
}

func (c *Calendar) Footer() string {
	return `

`
}

func (c *Calendar) openTR() string  { return "
\n" }
func (c *Calendar) closeTR() string { return "
\n" }

func (c *Calendar) entry(t time.Time) string {
	d := c.days[t]
	day := fmt.Sprintf("%02d", t.Day())
	class := fmt.Sprintf("\t", d)
	close := "	\n"
	href := ""
	switch d {
	case free:
		date := fmt.Sprintf("%4d-%02d-%02d", t.Year(), t.Month(), t.Day())
		href = fmt.Sprintf("%d", date, t.Day()) // BookingDate is defined on the page/form itself
		class = fmt.Sprintf("\t", d)
	case busy:
		href = day
	case past:
		href = day
	}
	s := class + href + close
	return s
}

func (c *Calendar) HTML() string {
	s := c.Header()
	s += c.html()
	s += c.Footer()
	return s
}

func (c *Calendar) sort() times {
	keys := times{}
	for k := range c.days {
		keys = append(keys, k)
	}

	sort.Sort(keys)
	return keys
}

func (c *Calendar) html() string {
	keys := c.sort()

	s := c.heading()
	i := 0
	for _, k := range keys {
		if i%7 == 0 {
			if i > 0 {
				s += c.closeTR()
			}
			s += c.openTR()
		}
		s += c.entry(k)
		i++
	}
	s += c.closeTR()
	return s
}

// New creates a new month calendar based on d, d must be in the form: YYYY-MM-DD.
// D can also be the empty string, then the current date is assumed.
func New(d string) (*Calendar, error) {
	date, now := time.Now(), time.Now()
	if d != "" {
		var err error
		date, err = time.Parse("2006-01-02", d)
		if err != nil {
			return nil, err
		}
	}

	cal := &Calendar{days: make(map[time.Time]Available), start: date}

	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
	first := time.Date(date.Year(), date.Month(), 1, 0, 0, 0, 0, time.UTC)
	last := time.Date(date.Year(), date.Month()+1, 1, 0, 0, 0, 0, time.UTC)
	last = last.Add(-24 * time.Hour)

	// Add the remaining days of the previous month.
	for i := 0; i < int(first.Weekday()); i++ {
		lastMonthDay := first.AddDate(0, 0, -1*(i+1))
		cal.days[lastMonthDay] = free

		if lastMonthDay.Before(today) {
			cal.days[lastMonthDay] = past
		}
	}

	// Loop from i to lastDay and add the entire month.
	for i := 1; i <= last.Day(); i++ {
		day := time.Date(date.Year(), date.Month(), i, 0, 0, 0, 0, time.UTC)

		cal.days[day] = free

		if day.Before(today) {
			cal.days[day] = past
		}
	}

	// These are dates in the new month.
	j := 1
	for i := int(last.Weekday()) + 1; i < 7; i++ {
		nextMonthDay := last.AddDate(0, 0, j)
		cal.days[nextMonthDay] = free

		if nextMonthDay.Before(today) {
			cal.days[nextMonthDay] = past
		}

		j++
	}
	times := cal.sort()
	if len(times) > 0 {
		cal.begin = times[0]
		cal.end = times[len(times)-1]
	}

	return cal, nil
}
<|endoftext|>


<commit_before>package main

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

	"sync"
)

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 uploadWorkers = uploadFlags.Int("workers", 4, "Number of upload workers")
var uploadRevs = uploadFlags.Int("revs", 0,
	"Number of old revisions to keep (-1 == all)")

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"
	}
	return def
}

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

	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]
	_, err = f.Seek(0, 0)
	if err != nil {
		return err
	}

	preq, err := http.NewRequest("PUT", dest, f)
	if err != nil {
		return err
	}
	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-Type", ctype)

	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
}

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

	listing, err := listStuff(baseUrl)
	for err != nil {
		return []string{}, err
	}
	for fn := range listing.Files {
		err = rmFile(baseUrl + "/" + fn)
		if err != nil {
			return []string{}, err
		}
	}
	children := make([]string, 0, len(listing.Dirs))
	for dn := range listing.Dirs {
		children = append(children, baseUrl+"/"+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:
				if req.remoteHash == "" {
					err = uploadFile(req.src, req.dest)
				} else {
					if localHash(req.src) != req.remoteHash {
						if *uploadVerbose {
							log.Printf("%v has changed, reupping",
								req.src)
						}
						err = uploadFile(req.src, req.dest)
					}
				}
			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)
				}
			} 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
	}

	retries := 3
	serverListing, err := listStuff(dest)
	for err != nil && retries > 0 {
		serverListing, err = listStuff(dest)
		time.Sleep(time.Second)
		retries--
	}
	if err != nil {
		// XXX: Stick an error somewhere
		return nil
	}

	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:
			localNames[c.Name()] = c
		}
	}

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

	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 + "/" + 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 + "/" + m, uploadFileOp, ""}
		}
	}

	if *uploadDelete && len(toRm) > 0 {
		for _, m := range toRm {
			ch <- uploadReq{"", dest + "/" + m, removeFileOp, ""}
			ch <- uploadReq{"", dest + "/" + 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() {
				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)

	if uploadFlags.NArg() < 1 {
		log.Fatalf("dest required")
	}

	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)
		if err != nil {
			log.Fatalf("Error uploading file: %v", err)
		}
	}
}
<commit_msg>Upload CSS as text/css<commit_after>package main

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

	"sync"
)

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 uploadWorkers = uploadFlags.Int("workers", 4, "Number of upload workers")
var uploadRevs = uploadFlags.Int("revs", 0,
	"Number of old revisions to keep (-1 == all)")

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 uploadFile(src, dest string) error {
	if *uploadVerbose {
		log.Printf("Uploading %v -> %v", src, dest)
	}

	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]
	_, err = f.Seek(0, 0)
	if err != nil {
		return err
	}

	preq, err := http.NewRequest("PUT", dest, f)
	if err != nil {
		return err
	}
	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-Type", ctype)

	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
}

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

	listing, err := listStuff(baseUrl)
	for err != nil {
		return []string{}, err
	}
	for fn := range listing.Files {
		err = rmFile(baseUrl + "/" + fn)
		if err != nil {
			return []string{}, err
		}
	}
	children := make([]string, 0, len(listing.Dirs))
	for dn := range listing.Dirs {
		children = append(children, baseUrl+"/"+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:
				if req.remoteHash == "" {
					err = uploadFile(req.src, req.dest)
				} else {
					if localHash(req.src) != req.remoteHash {
						if *uploadVerbose {
							log.Printf("%v has changed, reupping",
								req.src)
						}
						err = uploadFile(req.src, req.dest)
					}
				}
			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)
				}
			} 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
	}

	retries := 3
	serverListing, err := listStuff(dest)
	for err != nil && retries > 0 {
		serverListing, err = listStuff(dest)
		time.Sleep(time.Second)
		retries--
	}
	if err != nil {
		// XXX: Stick an error somewhere
		return nil
	}

	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:
			localNames[c.Name()] = c
		}
	}

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

	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 + "/" + 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 + "/" + m, uploadFileOp, ""}
		}
	}

	if *uploadDelete && len(toRm) > 0 {
		for _, m := range toRm {
			ch <- uploadReq{"", dest + "/" + m, removeFileOp, ""}
			ch <- uploadReq{"", dest + "/" + 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() {
				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)

	if uploadFlags.NArg() < 1 {
		log.Fatalf("dest required")
	}

	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)
		if err != nil {
			log.Fatalf("Error uploading file: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"testing"
	"time"
)

// Should serve a known static error page if all backend servers are down
// and object isn't in cache/stale.
// NB: ideally this should be a page that we control that has a mechanism
//     to alert us that it has been served.
func TestFailoverErrorPageAllServersDown(t *testing.T) {
	t.Error("Not implemented")
}

// Should serve a known static error page if all backend servers return a
// 5xx response and object isn't in cache/stale.
// NB: ideally this should be a page that we control that has a mechanism
//     to alert us that it has been served.
func TestFailoverErrorPageAllServers5xx(t *testing.T) {
	t.Error("Not implemented")
}

// Should back off requests against origin for a very short period of time
// if origin returns a 5xx response so as not to overwhelm it.
func TestFailoverOrigin5xxBackOff(t *testing.T) {
	t.Error("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin is down and
// object is beyond TTL but still in cache.
func TestFailoverOriginDownServeStale(t *testing.T) {
	t.Error("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin returns a 5xx
// response and object is beyond TTL but still in cache.
func TestFailoverOrigin5xxServeStale(t *testing.T) {
	const expectedResponseStale = "going off like stilton"
	const expectedResponseFresh = "as fresh as daisies"

	const respTTL = time.Duration(2 * time.Second)
	const respTTLWithBuffer = 5 * respTTL
	// Allow varnish's beresp.saintmode to expire.
	const waitSaintMode = time.Duration(5 * time.Second)
	headerValue := fmt.Sprintf("max-age=%.0f", respTTL.Seconds())

	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	url := fmt.Sprintf("https://%s/%s", *edgeHost, NewUUID())
	req, _ := http.NewRequest("GET", url, nil)

	for requestCount := 1; requestCount < 4; requestCount++ {
		var expectedBody string

		switch requestCount {
		case 1:
			expectedBody = expectedResponseStale
			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Cache-Control", headerValue)
				w.Write([]byte(expectedBody))
			})
		case 2:
			time.Sleep(respTTLWithBuffer)

			expectedBody = expectedResponseStale
			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusServiceUnavailable)
				w.Write([]byte(originServer.Name))
			})
		case 3:
			time.Sleep(waitSaintMode)

			expectedBody = expectedResponseFresh
			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(expectedBody))
			})
		}

		resp, err := client.RoundTrip(req)
		if err != nil {
			t.Fatal(err)
		}

		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Fatal(err)
		}
		if bodyStr := string(body); bodyStr != expectedBody {
			t.Errorf(
				"Request %d received incorrect response body. Expected %q, got %q",
				requestCount,
				expectedBody,
				bodyStr,
			)
		}
	}
}

// Should fallback to first mirror if origin is down and object is not in
// cache (active or stale).
func TestFailoverOriginDownUseFirstMirror(t *testing.T) {
	t.Error("Not implemented")
}

// Should fallback to first mirror if origin returns 5xx response and object
// is not in cache (active or stale).
func TestFailoverOrigin5xxUseFirstMirror(t *testing.T) {
	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received a request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	url := fmt.Sprintf("https://%s/%s", *edgeHost, NewUUID())
	req, _ := http.NewRequest("GET", url, nil)
	resp, err := client.RoundTrip(req)

	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to second mirror if both origin and first mirror are
// down.
func TestFailoverOriginDownFirstMirrorDownUseSecondMirror(t *testing.T) {
	t.Error("Not implemented")
}

// Should fallback to second mirror if both origin and first mirror return
// 5xx responses.
func TestFailoverOrigin5xxFirstMirror5xxUseSecondMirror(t *testing.T) {
	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})

	url := fmt.Sprintf("https://%s/%s", *edgeHost, NewUUID())
	req, _ := http.NewRequest("GET", url, nil)
	resp, err := client.RoundTrip(req)

	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should not fallback to mirror if origin returns a 5xx response with a
// No-Fallback header.
func TestFailoverNoFallbackHeader(t *testing.T) {
	t.Error("Not implemented")
}
<commit_msg>Make two additional requests when testing stale<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"testing"
	"time"
)

// Should serve a known static error page if all backend servers are down
// and object isn't in cache/stale.
// NB: ideally this should be a page that we control that has a mechanism
//     to alert us that it has been served.
func TestFailoverErrorPageAllServersDown(t *testing.T) {
	t.Error("Not implemented")
}

// Should serve a known static error page if all backend servers return a
// 5xx response and object isn't in cache/stale.
// NB: ideally this should be a page that we control that has a mechanism
//     to alert us that it has been served.
func TestFailoverErrorPageAllServers5xx(t *testing.T) {
	t.Error("Not implemented")
}

// Should back off requests against origin for a very short period of time
// if origin returns a 5xx response so as not to overwhelm it.
func TestFailoverOrigin5xxBackOff(t *testing.T) {
	t.Error("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin is down and
// object is beyond TTL but still in cache.
func TestFailoverOriginDownServeStale(t *testing.T) {
	t.Error("Not implemented")
}

// Should serve stale object and not hit mirror(s) if origin returns a 5xx
// response and object is beyond TTL but still in cache.
func TestFailoverOrigin5xxServeStale(t *testing.T) {
	const expectedResponseStale = "going off like stilton"
	const expectedResponseFresh = "as fresh as daisies"

	const respTTL = time.Duration(2 * time.Second)
	const respTTLWithBuffer = 5 * respTTL
	// Allow varnish's beresp.saintmode to expire.
	const waitSaintMode = time.Duration(5 * time.Second)
	headerValue := fmt.Sprintf("max-age=%.0f", respTTL.Seconds())

	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	url := fmt.Sprintf("https://%s/%s", *edgeHost, NewUUID())
	req, _ := http.NewRequest("GET", url, nil)

	var expectedBody string
	for requestCount := 1; requestCount < 6; requestCount++ {
		switch requestCount {
		case 1:
			expectedBody = expectedResponseStale

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Cache-Control", headerValue)
				w.Write([]byte(expectedBody))
			})
		case 2:
			time.Sleep(respTTLWithBuffer)
			expectedBody = expectedResponseStale

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusServiceUnavailable)
				w.Write([]byte(originServer.Name))
			})
		case 5:
			time.Sleep(waitSaintMode)
			expectedBody = expectedResponseFresh

			originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
				w.Write([]byte(expectedBody))
			})
		}

		resp, err := client.RoundTrip(req)
		if err != nil {
			t.Fatal(err)
		}

		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Fatal(err)
		}
		if bodyStr := string(body); bodyStr != expectedBody {
			t.Errorf(
				"Request %d received incorrect response body. Expected %q, got %q",
				requestCount,
				expectedBody,
				bodyStr,
			)
		}
	}
}

// Should fallback to first mirror if origin is down and object is not in
// cache (active or stale).
func TestFailoverOriginDownUseFirstMirror(t *testing.T) {
	t.Error("Not implemented")
}

// Should fallback to first mirror if origin returns 5xx response and object
// is not in cache (active or stale).
func TestFailoverOrigin5xxUseFirstMirror(t *testing.T) {
	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		t.Errorf("Server %s received a request and it shouldn't have", name)
		w.Write([]byte(name))
	})

	url := fmt.Sprintf("https://%s/%s", *edgeHost, NewUUID())
	req, _ := http.NewRequest("GET", url, nil)
	resp, err := client.RoundTrip(req)

	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should fallback to second mirror if both origin and first mirror are
// down.
func TestFailoverOriginDownFirstMirrorDownUseSecondMirror(t *testing.T) {
	t.Error("Not implemented")
}

// Should fallback to second mirror if both origin and first mirror return
// 5xx responses.
func TestFailoverOrigin5xxFirstMirror5xxUseSecondMirror(t *testing.T) {
	expectedBody := "lucky golden ticket"
	expectedStatus := http.StatusOK
	backendsSawRequest := map[string]bool{}

	originServer.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := originServer.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer1.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer1.Name
		if !backendsSawRequest[name] {
			w.WriteHeader(http.StatusServiceUnavailable)
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
		}
		w.Write([]byte(name))
	})
	backupServer2.SwitchHandler(func(w http.ResponseWriter, r *http.Request) {
		name := backupServer2.Name
		if !backendsSawRequest[name] {
			w.Write([]byte(expectedBody))
			backendsSawRequest[name] = true
		} else {
			t.Errorf("Server %s received more than one request", name)
			w.Write([]byte(name))
		}
	})

	url := fmt.Sprintf("https://%s/%s", *edgeHost, NewUUID())
	req, _ := http.NewRequest("GET", url, nil)
	resp, err := client.RoundTrip(req)

	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != expectedStatus {
		t.Errorf(
			"Received incorrect status code. Expected %d, got %d",
			expectedStatus,
			resp.StatusCode,
		)
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}
	if bodyStr := string(body); bodyStr != expectedBody {
		t.Errorf(
			"Received incorrect response body. Expected %q, got %q",
			expectedBody,
			bodyStr,
		)
	}
}

// Should not fallback to mirror if origin returns a 5xx response with a
// No-Fallback header.
func TestFailoverNoFallbackHeader(t *testing.T) {
	t.Error("Not implemented")
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"

	"github.com/kataras/iris"
)

var (
	errStaticDataNotFound = "File: %s was not found"
)

func IrisAPI() *iris.Framework {

	// TODO
	// iris.Config.Gzip = true

	api := iris.New()
	// define the api
	api.Post("/api/schemas/:id", func(ctx *iris.Context) {

		id := ctx.Param("id")

		// defaults
		jSONSchema := JSONSchema{
			Template:   "csv",
			EmptyValue: "",
			Count:      10,
		}

		parser := ValidatingJSONRequestParser{}
		template, err := parser.parse(ctx, &jSONSchema, id)
		if err != nil {
			emmitError(http.StatusBadRequest, err.Error(), ctx)
			return
		}

		// TODO slice to chunks + in separate go routines
		context := Context{
			FromIndex:    0,
			ToIndex:      jSONSchema.Count,
			CurrentIndex: []int{0},
		}
		ctx.SetContentType(template.ContentType(&context))

		out, err := template.Generate(&context)
		if err != nil {
			emmitError(http.StatusInternalServerError, err.Error(), ctx)
			return
		}

		println("POST /api/schemas/" + id)

		ctx.Write(out)
	})

	return api
}

func main() {
	IrisAPI().Listen(":9292")
}

func emmitError(errorCode int, errorString string, ctx *iris.Context) {
	ctx.EmitError(errorCode)
	ctx.Write(": " + errorString)
}
<commit_msg>feat: introduced chunked response<commit_after>package main

import (
	"bufio"
	"fmt"
	"net/http"

	"github.com/kataras/iris"
)

var (
	errStaticDataNotFound = "File: %s was not found"
)

func IrisAPI() *iris.Framework {

	api := iris.New()

	// TODO Gzip
	// iris.Config.Gzip = true

	// define the api
	api.Post("/api/schemas/:id", func(ctx *iris.Context) {

		id := ctx.Param("id")
		println("POST /api/schemas/" + id)

		// defaults
		jSONSchema := JSONSchema{
			Template:   "csv",
			EmptyValue: "",
			Count:      10,
		}

		parser := ValidatingJSONRequestParser{}
		template, err := parser.parse(ctx, &jSONSchema, id)
		if err != nil {
			emmitError(http.StatusBadRequest, err.Error(), ctx)
			return
		}

		// TODO slice to chunks + in separate go routines
		context := Context{
			FromIndex:    0,
			ToIndex:      jSONSchema.Count,
			CurrentIndex: []int{0},
		}
		ctx.SetContentType(template.ContentType(&context))
		ctx.StreamWriter(func(writer *bufio.Writer) {
			// TODO Gzip
			// ctx.Response.WriteGzip()

			out, err := template.Generate(&context)
			if err != nil {
				emmitError(http.StatusInternalServerError, err.Error(), ctx)
				return
			}

			fmt.Fprintf(writer, "%s", out)

			writer.Flush()
		})

	})

	return api
}

func main() {
	IrisAPI().Listen(":9292")
}

func emmitError(errorCode int, errorString string, ctx *iris.Context) {
	ctx.EmitError(errorCode)
	ctx.Write(": " + errorString)
}
<|endoftext|>