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

// Package auth provides a set of user authentication handlers for the ozzo routing package.
package jwt

import (
	"github.com/dgrijalva/jwt-go"
	"github.com/go-ozzo/ozzo-routing"
	"net/http"
	"strings"
        "log"
)

// User is the key used to store and retrieve the user identity information in routing.Context
const User = "User"

type JWTConfig struct {
	Alg     string
	Secret  string
	Expires int64
}

// Identity represents an authenticated user. If a user is successfully authenticated by
// an auth handler (Basic, Bearer, or Query), an Identity object will be made available for injection.
type Identity interface{}

type Payload interface{}
type JWTPayload map[string]interface{}

// DefaultRealm is the default realm name for HTTP authentication. It is used by HTTP authentication based on
// Basic and Bearer.
var DefaultRealm = "API"

// TokenAuthFunc is the function for authenticating a user based on a secret token.
type TokenAuthFunc func(c *routing.Context, payload JWTPayload) (Payload, error)

// JWT returns a routing.Handler that performs HTTP authentication based on bearer token.
// It can be used like the following:
//
// Example
//
// By default, the auth realm is named as "API". You may customize it by specifying the realm parameter.
//
// When authentication fails, a "WWW-Authenticate" header will be sent, and an http.StatusUnauthorized
// error will be reported via routing.Context.Error().

func JWT(fn TokenAuthFunc, jwtConfig JWTConfig, realm ...string) routing.Handler {
	name := DefaultRealm
	if len(realm) > 0 {
		name = realm[0]
	}
	return func(c *routing.Context) error {

		header := c.Request.Header.Get("Authorization")
		if !strings.HasPrefix(header, "Bearer ") {
			c.Response.Header().Set("WWW-Authenticate", `Bearer realm="`+name+`"`)
			return routing.NewHTTPError(http.StatusUnauthorized, "ivalid credential")
		}

		token, err := jwt.Parse(header[7:], func(t *jwt.Token) (interface{}, error) { return []byte(jwtConfig.Secret), nil })
                log.Println(token, err)
		if err != nil {
			c.Response.Header().Set("WWW-Authenticate", `Bearer realm="`+name+`"`)
			return routing.NewHTTPError(http.StatusUnauthorized, err.Error())
		}

		payload := token.Claims

		identity, e := fn(c, payload)
		if e == nil {
			c.Set(User, identity)
			return nil
		}

		c.Response.Header().Set("WWW-Authenticate", `Bearer realm="`+name+`"`)
		return routing.NewHTTPError(http.StatusUnauthorized, e.Error())

	}
}

// CreateToken returns token string
func CreateToken(jwtConfig JWTConfig, payload JWTPayload) (string, error) {
	signingMethod := jwt.SigningMethodHS256

	switch jwtConfig.Alg {
	case "HS384":
		signingMethod = jwt.SigningMethodHS384
	case "HS512":
		signingMethod = jwt.SigningMethodHS512
	}

	token := jwt.New(signingMethod)
	token.Claims = payload
	token.Claims["exp"] = jwtConfig.Expires

	tokenString, err := token.SignedString([]byte(jwtConfig.Secret))

	return tokenString, err

}
<commit_msg>Remove debug<commit_after>// Copyright 2016  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package auth provides a set of user authentication handlers for the ozzo routing package.
package jwt

import (
	"github.com/dgrijalva/jwt-go"
	"github.com/go-ozzo/ozzo-routing"
	"net/http"
	"strings"
)

// User is the key used to store and retrieve the user identity information in routing.Context
const User = "User"

type JWTConfig struct {
	Alg     string
	Secret  string
	Expires int64
}

// Identity represents an authenticated user. If a user is successfully authenticated by
// an auth handler (Basic, Bearer, or Query), an Identity object will be made available for injection.
type Identity interface{}

type Payload interface{}
type JWTPayload map[string]interface{}

// DefaultRealm is the default realm name for HTTP authentication. It is used by HTTP authentication based on
// Basic and Bearer.
var DefaultRealm = "API"

// TokenAuthFunc is the function for authenticating a user based on a secret token.
type TokenAuthFunc func(c *routing.Context, payload JWTPayload) (Payload, error)

// JWT returns a routing.Handler that performs HTTP authentication based on bearer token.
// It can be used like the following:
//
// Example
//
// By default, the auth realm is named as "API". You may customize it by specifying the realm parameter.
//
// When authentication fails, a "WWW-Authenticate" header will be sent, and an http.StatusUnauthorized
// error will be reported via routing.Context.Error().

func JWT(fn TokenAuthFunc, jwtConfig JWTConfig, realm ...string) routing.Handler {
	name := DefaultRealm
	if len(realm) > 0 {
		name = realm[0]
	}
	return func(c *routing.Context) error {

		header := c.Request.Header.Get("Authorization")
		if !strings.HasPrefix(header, "Bearer ") {
			c.Response.Header().Set("WWW-Authenticate", `Bearer realm="`+name+`"`)
			return routing.NewHTTPError(http.StatusUnauthorized, "ivalid credential")
		}

		token, err := jwt.Parse(header[7:], func(t *jwt.Token) (interface{}, error) { return []byte(jwtConfig.Secret), nil })
		if err != nil {
			c.Response.Header().Set("WWW-Authenticate", `Bearer realm="`+name+`"`)
			return routing.NewHTTPError(http.StatusUnauthorized, err.Error())
		}

		payload := token.Claims

		identity, e := fn(c, payload)
		if e == nil {
			c.Set(User, identity)
			return nil
		}

		c.Response.Header().Set("WWW-Authenticate", `Bearer realm="`+name+`"`)
		return routing.NewHTTPError(http.StatusUnauthorized, e.Error())

	}
}

// CreateToken returns token string
func CreateToken(jwtConfig JWTConfig, payload JWTPayload) (string, error) {
	signingMethod := jwt.SigningMethodHS256

	switch jwtConfig.Alg {
	case "HS384":
		signingMethod = jwt.SigningMethodHS384
	case "HS512":
		signingMethod = jwt.SigningMethodHS512
	}

	token := jwt.New(signingMethod)
	token.Claims = payload
	token.Claims["exp"] = jwtConfig.Expires

	tokenString, err := token.SignedString([]byte(jwtConfig.Secret))

	return tokenString, err

}
<|endoftext|>


<commit_before>
package main

import (
        "os"
        "fmt"
        "log"
        "time"
        "flag"
        "runtime"
        )

const VERSION = "0.1.1"

var (
    myId string
    cache *Cache
    runningCheckers = 0
)

func addCheck(line string) {
    check, err := NewCheck(line)
    if err != nil {
        log.Printf("Warning: got invalid data on the \"dead\" " +
            "channel (\"%s\")\n", line)
        return
    }
    if check.BackendGroupLength <= 1 {
        // Add the check only if the frontend is scaled to several
        // backends (backend is part of a group)
        return
    }
    // Set all the callbacks for the check. They will be called during
    // the PingUrl at several steps
    check.SetDeadCallback(func () {
        cache.MarkBackendDead(check)
    })
    check.SetAliveCallback(func () {
        cache.MarkBackendAlive(check)
    })
    check.SetCheckIfBreakCallback(func () bool {
        return cache.IsUnlockedBackend(check)
    })
    check.SetExitCallback(func () {
        runningCheckers -= 1
        cache.UnlockBackend(check)
    })
    r := cache.LockBackend(check)
    if r == false {
        // Backend already locked
        return
    }
    // Check the URL at a regular interval
    go check.PingUrl()
    runningCheckers += 1
    log.Printf("Added check for: %s\n", check.BackendUrl)
}

/*
 * Prints some stats on runtime
 */
func printStats() {
    go func () {
        for {
            log.Printf("%d backend URLs are being tested\n", runningCheckers)
            time.Sleep(time.Duration(1) * time.Minute)
        }
    }()
}

func parseFlags() {
    parseDuration := func (v *time.Duration, n string, def int, help string) {
        i := flag.Int(n, def, help)
        *v = time.Duration(*i) * time.Second
    }
    flag.StringVar(&httpMethod, "method", HTTP_METHOD,
        "HTTP method")
    parseDuration(&checkInterval, "interval", CHECK_INTERVAL,
        "Check interval (seconds)")
    parseDuration(&connectionTimeout, "connect", CONNECTION_TIMEOUT,
        "TCP connection timeout (seconds)")
    parseDuration(&ioTimeout, "io", IO_TIMEOUT,
        "Socket read/write timeout (seconds)")
    flag.StringVar(&userAgent, "agent", USER_AGENT,
        "HTTP User-Agent header")
    flag.StringVar(&redisAddress, "redis", REDIS_ADDRESS,
        "Network address of Redis")
    flag.Parse()
}

func main() {
    var (
        err error
        hostname string
        )
    for  _, arg := range os.Args {
        if !(arg == "-v" || arg == "--version" || arg == "-version") {
            continue
        }
        fmt.Println("hchecker version", VERSION)
        os.Exit(0)
    }
    parseFlags()
    runtime.GOMAXPROCS(runtime.NumCPU())
    hostname, _ = os.Hostname()
    myId = fmt.Sprintf("%s#%d", hostname, os.Getpid())
    // Prefix each line of log
    log.SetPrefix(myId + " ")
    cache, err = NewCache()
    if err != nil {
        log.Println(err.Error())
        os.Exit(1)
    }
    printStats()
    err = cache.ListenToChannel("dead", addCheck)
    if err != nil {
        log.Println(err.Error())
    }
}
<commit_msg>Bumped version to 0.1.2<commit_after>
package main

import (
        "os"
        "fmt"
        "log"
        "time"
        "flag"
        "runtime"
        )

const VERSION = "0.1.2"

var (
    myId string
    cache *Cache
    runningCheckers = 0
)

func addCheck(line string) {
    check, err := NewCheck(line)
    if err != nil {
        log.Printf("Warning: got invalid data on the \"dead\" " +
            "channel (\"%s\")\n", line)
        return
    }
    if check.BackendGroupLength <= 1 {
        // Add the check only if the frontend is scaled to several
        // backends (backend is part of a group)
        return
    }
    // Set all the callbacks for the check. They will be called during
    // the PingUrl at several steps
    check.SetDeadCallback(func () {
        cache.MarkBackendDead(check)
    })
    check.SetAliveCallback(func () {
        cache.MarkBackendAlive(check)
    })
    check.SetCheckIfBreakCallback(func () bool {
        return cache.IsUnlockedBackend(check)
    })
    check.SetExitCallback(func () {
        runningCheckers -= 1
        cache.UnlockBackend(check)
    })
    r := cache.LockBackend(check)
    if r == false {
        // Backend already locked
        return
    }
    // Check the URL at a regular interval
    go check.PingUrl()
    runningCheckers += 1
    log.Printf("Added check for: %s\n", check.BackendUrl)
}

/*
 * Prints some stats on runtime
 */
func printStats() {
    go func () {
        for {
            log.Printf("%d backend URLs are being tested\n", runningCheckers)
            time.Sleep(time.Duration(1) * time.Minute)
        }
    }()
}

func parseFlags() {
    parseDuration := func (v *time.Duration, n string, def int, help string) {
        i := flag.Int(n, def, help)
        *v = time.Duration(*i) * time.Second
    }
    flag.StringVar(&httpMethod, "method", HTTP_METHOD,
        "HTTP method")
    parseDuration(&checkInterval, "interval", CHECK_INTERVAL,
        "Check interval (seconds)")
    parseDuration(&connectionTimeout, "connect", CONNECTION_TIMEOUT,
        "TCP connection timeout (seconds)")
    parseDuration(&ioTimeout, "io", IO_TIMEOUT,
        "Socket read/write timeout (seconds)")
    flag.StringVar(&userAgent, "agent", USER_AGENT,
        "HTTP User-Agent header")
    flag.StringVar(&redisAddress, "redis", REDIS_ADDRESS,
        "Network address of Redis")
    flag.Parse()
}

func main() {
    var (
        err error
        hostname string
        )
    for  _, arg := range os.Args {
        if !(arg == "-v" || arg == "--version" || arg == "-version") {
            continue
        }
        fmt.Println("hchecker version", VERSION)
        os.Exit(0)
    }
    parseFlags()
    runtime.GOMAXPROCS(runtime.NumCPU())
    hostname, _ = os.Hostname()
    myId = fmt.Sprintf("%s#%d", hostname, os.Getpid())
    // Prefix each line of log
    log.SetPrefix(myId + " ")
    cache, err = NewCache()
    if err != nil {
        log.Println(err.Error())
        os.Exit(1)
    }
    printStats()
    err = cache.ListenToChannel("dead", addCheck)
    if err != nil {
        log.Println(err.Error())
    }
}
<|endoftext|>


<commit_before>package operator

import (
	"fmt"
	"sync"

	"github.com/giantswarm/azuretpr"
	"github.com/giantswarm/certificatetpr"
	"github.com/giantswarm/microerror"
	"github.com/giantswarm/micrologger"
	"github.com/giantswarm/operatorkit/tpr"
	"github.com/spf13/viper"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"

	"github.com/giantswarm/azure-operator/flag"
)

const (
	TPRVersion     = "v1"
	TPRDescription = "Managed Kubernetes clusters in Giantnetes running on Azure"
)

// Config represents the configuration used to create an Operator service.
type Config struct {
	// Dependencies.
	Logger    micrologger.Logger
	K8sClient kubernetes.Interface

	// Settings.
	Flag  *flag.Flag
	Viper *viper.Viper
}

// DefaultConfig provides a default configuration to create a new operator
// service by best effort.
func DefaultConfig() Config {
	return Config{
		// Dependencies.
		K8sClient: nil,
		Logger:    nil,

		// Settings.
		Flag:  nil,
		Viper: nil,
	}
}

// Service implements the Operator service interface.
type Service struct {
	Config

	// Internals.
	bootOnce sync.Once
	tpr      *tpr.TPR
}

// New creates a new configured Operator service.
func New(config Config) (*Service, error) {
	// Dependencies.
	if config.K8sClient == nil {
		return nil, microerror.Maskf(invalidConfigError, "kubernetes client must not be empty")
	}
	if config.Logger == nil {
		return nil, microerror.Maskf(invalidConfigError, "logger must not be empty")
	}

	// Settings.
	if config.Flag == nil {
		return nil, microerror.Maskf(invalidConfigError, "flag must not be empty")
	}
	if config.Viper == nil {
		return nil, microerror.Maskf(invalidConfigError, "viper must not be empty")
	}

	tprConfig := tpr.Config{
		K8sClient:   config.K8sClient,
		Logger:      config.Logger,
		Name:        certificatetpr.Name,
		Version:     TPRVersion,
		Description: TPRDescription,
	}
	tpr, err := tpr.New(tprConfig)
	if err != nil {
		return nil, microerror.Maskf(err, "creating TPR for %#v", tprConfig)
	}

	newService := &Service{
		Config: config,

		// Internals
		bootOnce: sync.Once{},
		tpr:      tpr,
	}

	return newService, nil
}

// Boot starts the service and implements the watch for azuretpr.
func (s *Service) Boot() {
	s.bootOnce.Do(func() {
		err := s.tpr.CreateAndWait()
		if tpr.IsAlreadyExists(err) {
			s.Logger.Log("debug", "third party resource already exists")
		} else if err != nil {
			s.Logger.Log("error", fmt.Sprintf("%#v", err))
			return
		}

		s.Logger.Log("debug", "starting list/watch")

		newResourceEventHandler := &cache.ResourceEventHandlerFuncs{
			AddFunc:    s.addFunc,
			DeleteFunc: s.deleteFunc,
		}
		newZeroObjectFactory := &tpr.ZeroObjectFactoryFuncs{
			NewObjectFunc:     func() runtime.Object { return &azuretpr.CustomObject{} },
			NewObjectListFunc: func() runtime.Object { return &azuretpr.List{} },
		}

		s.tpr.NewInformer(newResourceEventHandler, newZeroObjectFactory).Run(nil)
	})
}

// addFunc TODO
func (s *Service) addFunc(obj interface{}) {
	customObject := *obj.(*azuretpr.CustomObject)
	s.Config.Logger.Log("debug", fmt.Sprintf("creating custom object for cluster '%s'", customObject.Spec.Cluster.Cluster.ID))
}

// deleteFunc TODO
func (s *Service) deleteFunc(obj interface{}) {
	customObject := *obj.(*azuretpr.CustomObject)
	s.Config.Logger.Log("debug", fmt.Sprintf("deleting custom object for cluster '%s'", customObject.Spec.Cluster.Cluster.ID))
}
<commit_msg>Fix TPR name<commit_after>package operator

import (
	"fmt"
	"sync"

	"github.com/giantswarm/azuretpr"
	"github.com/giantswarm/microerror"
	"github.com/giantswarm/micrologger"
	"github.com/giantswarm/operatorkit/tpr"
	"github.com/spf13/viper"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"

	"github.com/giantswarm/azure-operator/flag"
)

const (
	TPRVersion     = "v1"
	TPRDescription = "Managed Kubernetes clusters in Giantnetes running on Azure"
)

// Config represents the configuration used to create an Operator service.
type Config struct {
	// Dependencies.
	Logger    micrologger.Logger
	K8sClient kubernetes.Interface

	// Settings.
	Flag  *flag.Flag
	Viper *viper.Viper
}

// DefaultConfig provides a default configuration to create a new operator
// service by best effort.
func DefaultConfig() Config {
	return Config{
		// Dependencies.
		K8sClient: nil,
		Logger:    nil,

		// Settings.
		Flag:  nil,
		Viper: nil,
	}
}

// Service implements the Operator service interface.
type Service struct {
	Config

	// Internals.
	bootOnce sync.Once
	tpr      *tpr.TPR
}

// New creates a new configured Operator service.
func New(config Config) (*Service, error) {
	// Dependencies.
	if config.K8sClient == nil {
		return nil, microerror.Maskf(invalidConfigError, "kubernetes client must not be empty")
	}
	if config.Logger == nil {
		return nil, microerror.Maskf(invalidConfigError, "logger must not be empty")
	}

	// Settings.
	if config.Flag == nil {
		return nil, microerror.Maskf(invalidConfigError, "flag must not be empty")
	}
	if config.Viper == nil {
		return nil, microerror.Maskf(invalidConfigError, "viper must not be empty")
	}

	tprConfig := tpr.Config{
		K8sClient:   config.K8sClient,
		Logger:      config.Logger,
		Name:        azuretpr.Name,
		Version:     TPRVersion,
		Description: TPRDescription,
	}
	tpr, err := tpr.New(tprConfig)
	if err != nil {
		return nil, microerror.Maskf(err, "creating TPR for %#v", tprConfig)
	}

	newService := &Service{
		Config: config,

		// Internals
		bootOnce: sync.Once{},
		tpr:      tpr,
	}

	return newService, nil
}

// Boot starts the service and implements the watch for azuretpr.
func (s *Service) Boot() {
	s.bootOnce.Do(func() {
		err := s.tpr.CreateAndWait()
		if tpr.IsAlreadyExists(err) {
			s.Logger.Log("debug", "third party resource already exists")
		} else if err != nil {
			s.Logger.Log("error", fmt.Sprintf("%#v", err))
			return
		}

		s.Logger.Log("debug", "starting list/watch")

		newResourceEventHandler := &cache.ResourceEventHandlerFuncs{
			AddFunc:    s.addFunc,
			DeleteFunc: s.deleteFunc,
		}
		newZeroObjectFactory := &tpr.ZeroObjectFactoryFuncs{
			NewObjectFunc:     func() runtime.Object { return &azuretpr.CustomObject{} },
			NewObjectListFunc: func() runtime.Object { return &azuretpr.List{} },
		}

		s.tpr.NewInformer(newResourceEventHandler, newZeroObjectFactory).Run(nil)
	})
}

// addFunc TODO
func (s *Service) addFunc(obj interface{}) {
	customObject := *obj.(*azuretpr.CustomObject)
	s.Config.Logger.Log("debug", fmt.Sprintf("creating custom object for cluster '%s'", customObject.Spec.Cluster.Cluster.ID))
}

// deleteFunc TODO
func (s *Service) deleteFunc(obj interface{}) {
	customObject := *obj.(*azuretpr.CustomObject)
	s.Config.Logger.Log("debug", fmt.Sprintf("deleting custom object for cluster '%s'", customObject.Spec.Cluster.Cluster.ID))
}
<|endoftext|>


<commit_before>package checkin

import (
	"io"
	"log"
	"os"
	"strings"

	"github.com/gin-gonic/gin"
	mod "github.com/jim3mar/tidy/models/checkin"
	//"github.com/jim3mar/tidy/models/user"
	"net/http"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	//"encoding/json"
	//"log"
	//"strconv"
	"time"
)

type CheckInResource struct {
	Mongo    *mgo.Session
	CollCI   *mgo.Collection
	CollUser *mgo.Collection
}

func (cr *CheckInResource) Init(session *mgo.Session) {
	cr.Mongo = session
	cr.CollCI = cr.Mongo.DB("tidy").C("checkin")
	cr.CollUser = cr.Mongo.DB("tidy").C("user")
}

// CheckIn do checkin task for special user id
// Method: POST
func (cr *CheckInResource) CheckIn(c *gin.Context) {
	now := time.Now()
	content := c.PostForm("content")
	uidString := c.PostForm("uid")
	img := c.PostForm("img")
	log.Print("Checkin user_id: " + uidString)
	uid := bson.ObjectIdHex(uidString)
	ciData := &mod.CheckIn{
		Id_:         bson.NewObjectId(),
		UserId:      uid,
		Content:     content,
		CreateAt:    now,
		CreateDay:   now.Day(),
		CreateMonth: int(now.Month()),
		CreateYear:  now.Year(),
		CreateHour:  now.Hour(),
		CreateMin:   now.Minute(),
		CreateSec:   now.Second(),
		Timestamp:   now.Unix(),
		Images:      []string{img},
	}
	err := cr.CollCI.Insert(ciData)
	if err != nil {
		panic(err)
	}
	err = cr.CollUser.Update(
		bson.M{
			"_id": uid,
		},
		bson.M{
			"$inc": bson.M{
				"continuous": 1,
			},
			"$set": bson.M{
				"last_checkin": *ciData,
			},
		})

	if err != nil {
		panic(err)
	}

	c.JSON(http.StatusOK, now.Unix())
}

// ListCheckIn return all checkin records
// Method: GET
func (cr *CheckInResource) ListCheckIn(c *gin.Context) {
	col := cr.Mongo.DB("tidy").C("checkin")
	uid := bson.ObjectIdHex(c.DefaultQuery("uid", ""))
	//objectId := bson.ObjectIdHex(id)
	//log.Print("user_id: " + uid)
	var ci []mod.CheckIn
	col.Find(bson.M{"user_id": uid}).All(&ci)
	//col.Find(nil).All(&ci)
	//log.Printf("%s", ci)
	c.JSON(http.StatusOK, ci)
}

func (cr *CheckInResource) UploadImg(c *gin.Context) {
	file, header, err := c.Request.FormFile("file")
	//file, _, err := c.Request.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, "")
		return
	}
	//log.Println(header)
	//filename split
	fns := strings.Split(header.Filename, ".")
	log.Println(fns)
	fileext := "png"
	if l := len(fns); l >= 2 {
		fileext = fns[l-1]
	}
	guid := bson.NewObjectId().Hex()
	filename := guid + "." + fileext
	log.Println(filename)
	//fmt.Println(header.Filename)
	out, err := os.Create("./tmp/" + filename)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()
	_, err = io.Copy(out, file)
	if err != nil {
		log.Fatal(err)
	}
	c.JSON(http.StatusOK, struct{
		GUID string `json:"guid"`
		Ext string `json:"ext"`
	}{
		GUID: guid,
		Ext: fileext,
	})
}
<commit_msg>support multi images<commit_after>package checkin

import (
	"io"
	"log"
	"os"
	"strings"

	"github.com/gin-gonic/gin"
	mod "github.com/jim3mar/tidy/models/checkin"
	//"github.com/jim3mar/tidy/models/user"
	"net/http"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	//"encoding/json"
	//"log"
	//"strconv"
	"time"
)

type CheckInResource struct {
	Mongo    *mgo.Session
	CollCI   *mgo.Collection
	CollUser *mgo.Collection
}

func (cr *CheckInResource) Init(session *mgo.Session) {
	cr.Mongo = session
	cr.CollCI = cr.Mongo.DB("tidy").C("checkin")
	cr.CollUser = cr.Mongo.DB("tidy").C("user")
}

// CheckIn do checkin task for special user id
// Method: POST
func (cr *CheckInResource) CheckIn(c *gin.Context) {
	now := time.Now()
	content := c.PostForm("content")
	uidString := c.PostForm("uid")
	img := c.PostForm("img")
	log.Print("Checkin user_id: " + uidString)
	uid := bson.ObjectIdHex(uidString)
	ciData := &mod.CheckIn{
		Id_:         bson.NewObjectId(),
		UserId:      uid,
		Content:     content,
		CreateAt:    now,
		CreateDay:   now.Day(),
		CreateMonth: int(now.Month()),
		CreateYear:  now.Year(),
		CreateHour:  now.Hour(),
		CreateMin:   now.Minute(),
		CreateSec:   now.Second(),
		Timestamp:   now.Unix(),
		Images:      strings.Split(img, "|"),
	}
	err := cr.CollCI.Insert(ciData)
	if err != nil {
		panic(err)
	}
	err = cr.CollUser.Update(
		bson.M{
			"_id": uid,
		},
		bson.M{
			"$inc": bson.M{
				"continuous": 1,
			},
			"$set": bson.M{
				"last_checkin": *ciData,
			},
		})

	if err != nil {
		panic(err)
	}

	c.JSON(http.StatusOK, now.Unix())
}

// ListCheckIn return all checkin records
// Method: GET
func (cr *CheckInResource) ListCheckIn(c *gin.Context) {
	col := cr.Mongo.DB("tidy").C("checkin")
	uid := bson.ObjectIdHex(c.DefaultQuery("uid", ""))
	//objectId := bson.ObjectIdHex(id)
	//log.Print("user_id: " + uid)
	var ci []mod.CheckIn
	col.Find(bson.M{"user_id": uid}).All(&ci)
	//col.Find(nil).All(&ci)
	//log.Printf("%s", ci)
	c.JSON(http.StatusOK, ci)
}

func (cr *CheckInResource) UploadImg(c *gin.Context) {
	file, header, err := c.Request.FormFile("file")
	//file, _, err := c.Request.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, "")
		return
	}
	//log.Println(header)
	//filename split
	fns := strings.Split(header.Filename, ".")
	log.Println(fns)
	fileext := "png"
	if l := len(fns); l >= 2 {
		fileext = fns[l-1]
	}
	guid := bson.NewObjectId().Hex()
	filename := guid + "." + fileext
	log.Println(filename)
	//fmt.Println(header.Filename)
	out, err := os.Create("./tmp/" + filename)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()
	_, err = io.Copy(out, file)
	if err != nil {
		log.Fatal(err)
	}
	c.JSON(http.StatusOK, struct {
		GUID string `json:"guid"`
		Ext  string `json:"ext"`
	}{
		GUID: guid,
		Ext:  fileext,
	})
}
<|endoftext|>


<commit_before>package scraper

import (
	"fmt"
	"log"
	"math"
	"sync"

	"github.com/influxdata/influxdb/models"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/storage"
)

var (
	_ Registry               = &Service{}
	_ storage.SampleAppender = &Service{}
)

// Service represents the scraper manager
type Service struct {
	PointsWriter interface {
		WritePoints(database, retentionPolicy string, consistencyLevel models.ConsistencyLevel, points []models.Point) error
	}
	mu sync.Mutex
	wg sync.WaitGroup

	open     bool
	running  bool
	closing  chan struct{}
	updating chan *config.Config

	configs []Config

	logger *log.Logger

	discoverers []Discoverer

	// TargetManager represents a scraping/discovery manager
	mgr interface {
		ApplyConfig(cfg *config.Config) error
		Stop()
		Start()
		Wait()
	}
}

// NewService creates a new scraper service
func NewService(c []Config, l *log.Logger) *Service {
	s := &Service{
		configs: c,
		logger:  l,
	}
	s.mgr = retrieval.NewTargetManager(s, NewLogger(l))
	return s
}

// Open starts the scraper service
func (s *Service) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.open {
		return nil
	}

	s.open = true
	s.updating = make(chan *config.Config)
	s.closing = make(chan struct{})

	go s.scrape()
	s.logger.Println("I! opened service")

	s.open = true
	return nil
}

// Close stops the scraper service
func (s *Service) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.open {
		return nil
	}

	s.open = false
	close(s.closing)

	s.wg.Wait()

	s.logger.Println("I! closed service")
	return nil
}

func (s *Service) scrape() {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		s.mu.Lock()
		pairs := s.pairs()
		conf := s.prom(pairs)
		s.mu.Unlock()

		s.mgr.ApplyConfig(conf)

		s.mu.Lock()
		// Need to check open if service was stopped just before this lock
		if s.open {
			s.mgr.Start()
			s.running = true
		}
		s.mu.Unlock()
		// Wait will block until mgr.Stop has been called.
		s.mgr.Wait()
	}()

	for {
		select {
		case <-s.closing:
			s.mu.Lock()
			// Need to check if the targetmanager is even running before Stopping
			if s.running {
				s.mgr.Stop()
				s.running = false
			}
			s.mu.Unlock()
			return
		case conf := <-s.updating:
			s.logger.Println("I! updating scraper service")
			s.mgr.ApplyConfig(conf)
		}
	}
}

// Append tranforms prometheus samples and inserts data into the tasks pipeline
func (s *Service) Append(sample *model.Sample) error {
	value := float64(sample.Value)
	// Remove all NaN values
	// TODO: Add counter stat for this variable
	if math.IsNaN(value) {
		return nil
	}

	var err error
	db := ""
	rp := ""
	job := ""

	tags := map[string]string{}
	for name, value := range sample.Metric {
		if name == "job" {
			db, rp, job, err = decodeJobName(string(value))
			continue
		}
		tags[string(name)] = string(value)
	}

	blacklist := func() bool {
		s.mu.Lock()
		defer s.mu.Unlock()
		// If instance is blacklisted then do not send to PointsWriter
		if instance, ok := tags["instance"]; ok {
			for i := range s.configs {
				if s.configs[i].Name == job {
					for _, listed := range s.configs[i].Blacklist {
						if instance == listed {
							return true
						}
					}
				}
			}
		}
		return false
	}
	if blacklist() {
		return nil
	}

	fields := models.Fields{
		"value": value,
	}

	pt, err := models.NewPoint(job, models.NewTags(tags), fields, sample.Timestamp.Time())
	if err != nil {
		return err
	}

	return s.PointsWriter.WritePoints(db, rp, models.ConsistencyLevelAny, []models.Point{pt})
}

// NeedsThrottling conforms to SampleAppender and never returns true currently.
func (s *Service) NeedsThrottling() bool {
	return false
}

type testOptions struct {
	Name string `json:"name"`
}

// TestOptions returns options that are allowed for the Test
func (s *Service) TestOptions() interface{} {
	return &testOptions{}
}

// Test tests the options for the scrapers
func (s *Service) Test(options interface{}) error {
	return nil
}

// Update will replace all scraper configurations and apply the configuration
// to the target manager
func (s *Service) Update(newConfigs []interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	configs := make([]Config, len(newConfigs))
	for i, c := range newConfigs {
		if config, ok := c.(Config); ok {
			configs[i] = config
		} else {
			return fmt.Errorf("unexpected config object type, got %T exp %T", c, config)
		}
	}

	s.configs = configs
	if s.open {
		pairs := s.pairs()
		conf := s.prom(pairs)
		select {
		case <-s.closing:
			return fmt.Errorf("error writing configuration to closed scraper")
		case s.updating <- conf:
		}
	}
	return nil
}

// prom assumes service is locked
func (s *Service) prom(pairs []Pair) *config.Config {
	conf := &config.Config{
		ScrapeConfigs: make([]*config.ScrapeConfig, len(pairs)),
	}

	for i, pair := range pairs {
		sc := pair.Scraper.Prom()
		pair.Discoverer.Prom(sc)
		conf.ScrapeConfigs[i] = sc
	}
	return conf
}

// Commit applies the configuration to the scraper
func (s *Service) Commit() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return nil
	}

	pairs := s.pairs()
	conf := s.prom(pairs)
	select {
	case <-s.closing:
		return fmt.Errorf("error writing configuration to closed registry")
	case s.updating <- conf:
	}

	return nil
}

// AddDiscoverer adds discoverer to the registry
func (s *Service) AddDiscoverer(discoverer Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	s.discoverers = append(s.discoverers, discoverer)
}

// RemoveDiscoverer removes discoverer from the registry
func (s *Service) RemoveDiscoverer(rm Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	for i, d := range s.discoverers {
		if d.ServiceID() == rm.ServiceID() && d.Service() == rm.Service() {
			s.discoverers = append(s.discoverers[:i], s.discoverers[i+1:]...)
		}
	}
}

// AddScrapers adds scrapers to the registry
func (s *Service) AddScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	s.configs = append(s.configs, scrapers...)
}

// RemoveScrapers removes scrapers from the registry
func (s *Service) RemoveScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	for _, rm := range scrapers {
		for i, c := range s.configs {
			if c.Name == rm.Name {
				s.configs = append(s.configs[:i], s.configs[i+1:]...)
			}
		}
	}
}

// pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) pairs() []Pair {
	pairs := []Pair{}
	for _, scr := range s.configs {
		if !scr.Enabled {
			continue
		}
		for _, d := range s.discoverers {
			if scr.DiscoverID == d.ServiceID() && scr.DiscoverService == d.Service() {
				pairs = append(pairs, Pair{
					Scraper:    scr,
					Discoverer: d,
				})
			}
		}
	}
	return pairs
}

// Pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) Pairs() []Pair {
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.pairs()
}
<commit_msg>fix panic on remove discoverer during shutdown process<commit_after>package scraper

import (
	"fmt"
	"log"
	"math"
	"sync"

	"github.com/influxdata/influxdb/models"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/storage"
)

var (
	_ Registry               = &Service{}
	_ storage.SampleAppender = &Service{}
)

// Service represents the scraper manager
type Service struct {
	PointsWriter interface {
		WritePoints(database, retentionPolicy string, consistencyLevel models.ConsistencyLevel, points []models.Point) error
	}
	mu sync.Mutex
	wg sync.WaitGroup

	open     bool
	running  bool
	closing  chan struct{}
	updating chan *config.Config

	configs []Config

	logger *log.Logger

	discoverers []Discoverer

	// TargetManager represents a scraping/discovery manager
	mgr interface {
		ApplyConfig(cfg *config.Config) error
		Stop()
		Start()
		Wait()
	}
}

// NewService creates a new scraper service
func NewService(c []Config, l *log.Logger) *Service {
	s := &Service{
		configs: c,
		logger:  l,
	}
	s.mgr = retrieval.NewTargetManager(s, NewLogger(l))
	return s
}

// Open starts the scraper service
func (s *Service) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.open {
		return nil
	}

	s.open = true
	s.updating = make(chan *config.Config)
	s.closing = make(chan struct{})

	go s.scrape()
	s.logger.Println("I! opened service")

	s.open = true
	return nil
}

// Close stops the scraper service
func (s *Service) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.open {
		return nil
	}

	s.open = false
	close(s.closing)

	s.wg.Wait()

	s.logger.Println("I! closed service")
	return nil
}

func (s *Service) scrape() {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		s.mu.Lock()
		pairs := s.pairs()
		conf := s.prom(pairs)
		s.mu.Unlock()

		s.mgr.ApplyConfig(conf)

		s.mu.Lock()
		// Need to check open if service was stopped just before this lock
		if s.open {
			s.mgr.Start()
			s.running = true
		}
		s.mu.Unlock()
		// Wait will block until mgr.Stop has been called.
		s.mgr.Wait()
	}()

	for {
		select {
		case <-s.closing:
			s.mu.Lock()
			// Need to check if the targetmanager is even running before Stopping
			if s.running {
				s.mgr.Stop()
				s.running = false
			}
			s.mu.Unlock()
			return
		case conf := <-s.updating:
			s.logger.Println("I! updating scraper service")
			s.mgr.ApplyConfig(conf)
		}
	}
}

// Append tranforms prometheus samples and inserts data into the tasks pipeline
func (s *Service) Append(sample *model.Sample) error {
	value := float64(sample.Value)
	// Remove all NaN values
	// TODO: Add counter stat for this variable
	if math.IsNaN(value) {
		return nil
	}

	var err error
	db := ""
	rp := ""
	job := ""

	tags := map[string]string{}
	for name, value := range sample.Metric {
		if name == "job" {
			db, rp, job, err = decodeJobName(string(value))
			continue
		}
		tags[string(name)] = string(value)
	}

	blacklist := func() bool {
		s.mu.Lock()
		defer s.mu.Unlock()
		// If instance is blacklisted then do not send to PointsWriter
		if instance, ok := tags["instance"]; ok {
			for i := range s.configs {
				if s.configs[i].Name == job {
					for _, listed := range s.configs[i].Blacklist {
						if instance == listed {
							return true
						}
					}
				}
			}
		}
		return false
	}
	if blacklist() {
		return nil
	}

	fields := models.Fields{
		"value": value,
	}

	pt, err := models.NewPoint(job, models.NewTags(tags), fields, sample.Timestamp.Time())
	if err != nil {
		return err
	}

	return s.PointsWriter.WritePoints(db, rp, models.ConsistencyLevelAny, []models.Point{pt})
}

// NeedsThrottling conforms to SampleAppender and never returns true currently.
func (s *Service) NeedsThrottling() bool {
	return false
}

type testOptions struct {
	Name string `json:"name"`
}

// TestOptions returns options that are allowed for the Test
func (s *Service) TestOptions() interface{} {
	return &testOptions{}
}

// Test tests the options for the scrapers
func (s *Service) Test(options interface{}) error {
	return nil
}

// Update will replace all scraper configurations and apply the configuration
// to the target manager
func (s *Service) Update(newConfigs []interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	configs := make([]Config, len(newConfigs))
	for i, c := range newConfigs {
		if config, ok := c.(Config); ok {
			configs[i] = config
		} else {
			return fmt.Errorf("unexpected config object type, got %T exp %T", c, config)
		}
	}

	s.configs = configs
	if s.open {
		pairs := s.pairs()
		conf := s.prom(pairs)
		select {
		case <-s.closing:
			return fmt.Errorf("error writing configuration to closed scraper")
		case s.updating <- conf:
		}
	}
	return nil
}

// prom assumes service is locked
func (s *Service) prom(pairs []Pair) *config.Config {
	conf := &config.Config{
		ScrapeConfigs: make([]*config.ScrapeConfig, len(pairs)),
	}

	for i, pair := range pairs {
		sc := pair.Scraper.Prom()
		pair.Discoverer.Prom(sc)
		conf.ScrapeConfigs[i] = sc
	}
	return conf
}

// Commit applies the configuration to the scraper
func (s *Service) Commit() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return nil
	}

	pairs := s.pairs()
	conf := s.prom(pairs)
	select {
	case <-s.closing:
		return fmt.Errorf("error writing configuration to closed registry")
	case s.updating <- conf:
	}

	return nil
}

// AddDiscoverer adds discoverer to the registry
func (s *Service) AddDiscoverer(discoverer Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	s.discoverers = append(s.discoverers, discoverer)
}

// RemoveDiscoverer removes discoverer from the registry
func (s *Service) RemoveDiscoverer(rm Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	filtered := s.discoverers[0:0]
	for _, d := range s.discoverers {
		if d.ServiceID() != rm.ServiceID() || d.Service() != rm.Service() {
			filtered = append(filtered, d)
		}
	}
	s.discoverers = filtered
}

// AddScrapers adds scrapers to the registry
func (s *Service) AddScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	s.configs = append(s.configs, scrapers...)
}

// RemoveScrapers removes scrapers from the registry
func (s *Service) RemoveScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	for _, rm := range scrapers {
		for i, c := range s.configs {
			if c.Name == rm.Name {
				s.configs = append(s.configs[:i], s.configs[i+1:]...)
			}
		}
	}
}

// pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) pairs() []Pair {
	pairs := []Pair{}
	for _, scr := range s.configs {
		if !scr.Enabled {
			continue
		}
		for _, d := range s.discoverers {
			if scr.DiscoverID == d.ServiceID() && scr.DiscoverService == d.Service() {
				pairs = append(pairs, Pair{
					Scraper:    scr,
					Discoverer: d,
				})
			}
		}
	}
	return pairs
}

// Pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) Pairs() []Pair {
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.pairs()
}
<|endoftext|>


<commit_before>package main

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

	"bufio"
	"flag"
	"fmt"
	"net"
	"net/http"
	"regexp"
)

// VirtualizeMode - default mode when Hoverfly looks for captured requests to respond
const VirtualizeMode = "virtualize"

// SynthesizeMode - all requests are sent to middleware to create response
const SynthesizeMode = "synthesize"

// ModifyMode - middleware is applied to outgoing and incoming traffic
const ModifyMode = "modify"

// CaptureMode - requests are captured and stored in cache
const CaptureMode = "capture"

// orPanic - wrapper for logging errors
func orPanic(err error) {
	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Panic("Got error.")
	}
}

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

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "should proxy capture requests")
	synthesize := flag.Bool("synthesize", false, "should proxy capture requests")
	modify := flag.Bool("modify", false, "should proxy only modify requests")

	destination := flag.String("destination", ".", "destination URI to catch")
	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	// metrics
	metrics := flag.Bool("metrics", false, "supply -metrics flag to enable metrics logging to stdout")

	// development
	dev := flag.Bool("dev", false, "supply -dev flag to serve directly from ./static/dist instead from statik binary")

	flag.Parse()

	// getting settings
	cfg := InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.verbose = *verbose

	if *dev {
		// making text pretty
		log.SetFormatter(&log.TextFormatter{})
	}

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.proxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.adminPort = *adminPort
	}

	// development settings
	cfg.development = *dev

	// overriding default middleware setting
	cfg.middleware = *middleware

	// setting default mode
	mode := VirtualizeMode

	if *capture {
		mode = CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = SynthesizeMode

		if cfg.middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = ModifyMode

		if cfg.middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// overriding default settings
	cfg.mode = mode

	// overriding destination
	cfg.destination = *destination

	proxy, dbClient := getNewHoverfly(cfg)
	defer dbClient.cache.db.Close()

	// starting admin interface
	go dbClient.startAdminInterface()

	// start metrics registry flush
	if *metrics {
		go dbClient.counter.Init()
	}

	log.Warn(http.ListenAndServe(fmt.Sprintf(":%s", cfg.proxyPort), proxy))
}

// getNewHoverfly returns a configured ProxyHttpServer and DBClient, also starts admin interface on configured port
func getNewHoverfly(cfg *Configuration) (*goproxy.ProxyHttpServer, DBClient) {

	// getting boltDB
	db := getDB(cfg.databaseName)

	cache := Cache{
		db:             db,
		requestsBucket: []byte(requestsBucketName),
	}

	counter := NewModeCounter()

	// getting connections
	d := DBClient{
		cache:   cache,
		http:    &http.Client{},
		cfg:     cfg,
		counter: counter,
	}

	// creating proxy
	proxy := goproxy.NewProxyHttpServer()

	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.cfg.destination))).
		HandleConnect(goproxy.AlwaysMitm)

	// enable curl -p for all hosts on port 80
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.cfg.destination))).
		HijackConnect(func(req *http.Request, client net.Conn, ctx *goproxy.ProxyCtx) {
		defer func() {
			if e := recover(); e != nil {
				ctx.Logf("error connecting to remote: %v", e)
				client.Write([]byte("HTTP/1.1 500 Cannot reach destination\r\n\r\n"))
			}
			client.Close()
		}()
		clientBuf := bufio.NewReadWriter(bufio.NewReader(client), bufio.NewWriter(client))
		remote, err := net.Dial("tcp", req.URL.Host)
		orPanic(err)
		remoteBuf := bufio.NewReadWriter(bufio.NewReader(remote), bufio.NewWriter(remote))
		for {
			req, err := http.ReadRequest(clientBuf.Reader)
			orPanic(err)
			orPanic(req.Write(remoteBuf))
			orPanic(remoteBuf.Flush())
			resp, err := http.ReadResponse(remoteBuf.Reader, req)

			orPanic(err)
			orPanic(resp.Write(clientBuf.Writer))
			orPanic(clientBuf.Flush())
		}
	})

	// processing connections
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(cfg.destination))).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			return d.processRequest(r)
		})

	// intercepts response
	proxy.OnResponse(goproxy.ReqHostMatches(regexp.MustCompile(cfg.destination))).DoFunc(
		func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			d.counter.Count(d.cfg.GetMode())
			return resp
		})

	proxy.Verbose = d.cfg.verbose
	// proxy starting message
	log.WithFields(log.Fields{
		"Destination": d.cfg.destination,
		"ProxyPort":   d.cfg.proxyPort,
		"Mode":        d.cfg.GetMode(),
	}).Info("Proxy prepared...")

	return proxy, d
}

func hoverflyError(req *http.Request, err error, msg string, statusCode int) *http.Response {
	return goproxy.NewResponse(req,
		goproxy.ContentTypeText, statusCode,
		fmt.Sprintf("Hoverfly Error! %s. Got error: %s \n", msg, err.Error()))
}

// processRequest - processes incoming requests and based on proxy state (record/playback)
// returns HTTP response.
func (d *DBClient) processRequest(req *http.Request) (*http.Request, *http.Response) {

	mode := d.cfg.GetMode()

	if mode == CaptureMode {
		newResponse, err := d.captureRequest(req)

		if err != nil {
			return req, hoverflyError(req, err, "Could not capture request", http.StatusServiceUnavailable)
		}
		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.cfg.middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("request and response captured")

		return req, newResponse

	} else if mode == SynthesizeMode {
		response, err := synthesizeResponse(req, d.cfg.middleware)

		if err != nil {
			return req, hoverflyError(req, err, "Could not create synthetic response!", http.StatusServiceUnavailable)
		}

		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.cfg.middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("synthetic response created successfuly")

		return req, response

	} else if mode == ModifyMode {
		response, err := d.modifyRequestResponse(req, d.cfg.middleware)

		if err != nil {
			log.WithFields(log.Fields{
				"error":      err.Error(),
				"middleware": d.cfg.middleware,
			}).Error("Got error when performing request modification")
			return req, hoverflyError(
				req,
				err,
				fmt.Sprintf("Middleware (%s) failed or something else happened!", d.cfg.middleware),
				http.StatusServiceUnavailable)
		}
		// returning modified response
		return req, response
	}

	newResponse := d.getResponse(req)
	return req, newResponse

}
<commit_msg>removed main function, moved to cmd/hoverfly/main.go, minor refactoring to use exported fields<commit_after>package hoverfly

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

	"bufio"
	"fmt"
	"net"
	"net/http"
	"regexp"
)

// VirtualizeMode - default mode when Hoverfly looks for captured requests to respond
const VirtualizeMode = "virtualize"

// SynthesizeMode - all requests are sent to middleware to create response
const SynthesizeMode = "synthesize"

// ModifyMode - middleware is applied to outgoing and incoming traffic
const ModifyMode = "modify"

// CaptureMode - requests are captured and stored in cache
const CaptureMode = "capture"

// orPanic - wrapper for logging errors
func orPanic(err error) {
	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Panic("Got error.")
	}
}

// GetNewHoverfly returns a configured ProxyHttpServer and DBClient
func GetNewHoverfly(cfg *Configuration) (*goproxy.ProxyHttpServer, DBClient) {

	// getting boltDB
	db := getDB(cfg.DatabaseName)

	cache := Cache{
		DS:             db,
		RequestsBucket: []byte(requestsBucketName),
	}

	counter := NewModeCounter()

	// getting connections
	d := DBClient{
		Cache:   cache,
		HTTP:    &http.Client{},
		Cfg:     cfg,
		Counter: counter,
	}

	// creating proxy
	proxy := goproxy.NewProxyHttpServer()

	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.Cfg.Destination))).
		HandleConnect(goproxy.AlwaysMitm)

	// enable curl -p for all hosts on port 80
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.Cfg.Destination))).
		HijackConnect(func(req *http.Request, client net.Conn, ctx *goproxy.ProxyCtx) {
		defer func() {
			if e := recover(); e != nil {
				ctx.Logf("error connecting to remote: %v", e)
				client.Write([]byte("HTTP/1.1 500 Cannot reach destination\r\n\r\n"))
			}
			client.Close()
		}()
		clientBuf := bufio.NewReadWriter(bufio.NewReader(client), bufio.NewWriter(client))
		remote, err := net.Dial("tcp", req.URL.Host)
		orPanic(err)
		remoteBuf := bufio.NewReadWriter(bufio.NewReader(remote), bufio.NewWriter(remote))
		for {
			req, err := http.ReadRequest(clientBuf.Reader)
			orPanic(err)
			orPanic(req.Write(remoteBuf))
			orPanic(remoteBuf.Flush())
			resp, err := http.ReadResponse(remoteBuf.Reader, req)

			orPanic(err)
			orPanic(resp.Write(clientBuf.Writer))
			orPanic(clientBuf.Flush())
		}
	})

	// processing connections
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(cfg.Destination))).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			return d.processRequest(r)
		})

	// intercepts response
	proxy.OnResponse(goproxy.ReqHostMatches(regexp.MustCompile(cfg.Destination))).DoFunc(
		func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			d.Counter.Count(d.Cfg.GetMode())
			return resp
		})

	proxy.Verbose = d.Cfg.Verbose
	// proxy starting message
	log.WithFields(log.Fields{
		"Destination": d.Cfg.Destination,
		"ProxyPort":   d.Cfg.ProxyPort,
		"Mode":        d.Cfg.GetMode(),
	}).Info("Proxy prepared...")

	return proxy, d
}

func hoverflyError(req *http.Request, err error, msg string, statusCode int) *http.Response {
	return goproxy.NewResponse(req,
		goproxy.ContentTypeText, statusCode,
		fmt.Sprintf("Hoverfly Error! %s. Got error: %s \n", msg, err.Error()))
}

// processRequest - processes incoming requests and based on proxy state (record/playback)
// returns HTTP response.
func (d *DBClient) processRequest(req *http.Request) (*http.Request, *http.Response) {

	mode := d.Cfg.GetMode()

	if mode == CaptureMode {
		newResponse, err := d.captureRequest(req)

		if err != nil {
			return req, hoverflyError(req, err, "Could not capture request", http.StatusServiceUnavailable)
		}
		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.Cfg.Middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("request and response captured")

		return req, newResponse

	} else if mode == SynthesizeMode {
		response, err := synthesizeResponse(req, d.Cfg.Middleware)

		if err != nil {
			return req, hoverflyError(req, err, "Could not create synthetic response!", http.StatusServiceUnavailable)
		}

		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.Cfg.Middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("synthetic response created successfuly")

		return req, response

	} else if mode == ModifyMode {
		response, err := d.modifyRequestResponse(req, d.Cfg.Middleware)

		if err != nil {
			log.WithFields(log.Fields{
				"error":      err.Error(),
				"middleware": d.Cfg.Middleware,
			}).Error("Got error when performing request modification")
			return req, hoverflyError(
				req,
				err,
				fmt.Sprintf("Middleware (%s) failed or something else happened!", d.Cfg.Middleware),
				http.StatusServiceUnavailable)
		}
		// returning modified response
		return req, response
	}

	newResponse := d.getResponse(req)
	return req, newResponse

}
<|endoftext|>


<commit_before>/*
   Copyright 2015 The Httpgzip Authors. See the AUTHORS file at the
   top-level directory of this distribution and at
   <https://xi2.org/x/httpgzip/m/AUTHORS>.

   This file is part of Httpgzip.

   Httpgzip is free software: you can redistribute it and/or modify it
   under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   Httpgzip is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with Httpgzip.  If not, see <https://www.gnu.org/licenses/>.
*/

// Package httpgzip implements an http.Handler wrapper adding gzip
// compression for appropriate requests.
//
// It attempts to properly parse the request's Accept-Encoding header
// according to RFC 2616 and does not do a simple string search for
// "gzip" (which will fail to do the correct thing for values such as
// "*" or "identity,gzip;q=0"). It will serve either gzip or identity
// content codings, or return 406 Not Acceptable status if it can do
// neither.
//
// It works correctly with handlers which honour Range request headers
// (such as http.FileServer) by removing the Range header for requests
// which prefer gzip encoding. This is necessary since Range requests
// apply to the gzipped content but the wrapped handler is not aware
// of the compression when it writes byte ranges. The Accept-Ranges
// header is also stripped from corresponding responses.
//
// A Content-Type header is set using http.DetectContentType if it is
// not set by the wrapped handler.
package httpgzip // import "xi2.org/x/httpgzip"

import (
	"bytes"
	"compress/gzip"
	"mime"
	"net/http"
	"strconv"
	"strings"
	"sync"
)

// DefaultContentTypes is the default list of content types for which
// a Handler considers gzip compression. This list originates from the
// file compression.conf within the Apache configuration found at
// https://html5boilerplate.com/.
var DefaultContentTypes = []string{
	"application/atom+xml",
	"application/javascript",
	"application/json",
	"application/ld+json",
	"application/manifest+json",
	"application/rdf+xml",
	"application/rss+xml",
	"application/schema+json",
	"application/vnd.geo+json",
	"application/vnd.ms-fontobject",
	"application/x-font-ttf",
	"application/x-javascript",
	"application/x-web-app-manifest+json",
	"application/xhtml+xml",
	"application/xml",
	"font/eot",
	"font/opentype",
	"image/bmp",
	"image/svg+xml",
	"image/vnd.microsoft.icon",
	"image/x-icon",
	"text/cache-manifest",
	"text/css",
	"text/html",
	"text/javascript",
	"text/plain",
	"text/vcard",
	"text/vnd.rim.location.xloc",
	"text/vtt",
	"text/x-component",
	"text/x-cross-domain-policy",
	"text/xml",
}

var gzipWriterPool = sync.Pool{
	New: func() interface{} { return gzip.NewWriter(nil) },
}

var gzipBufPool = sync.Pool{
	New: func() interface{} { return new(bytes.Buffer) },
}

// A gzipResponseWriter is a modified http.ResponseWriter. It adds
// gzip compression to certain responses, and there are two cases
// where this is done. Case 1 is when encs only allows gzip encoding
// and forbids identity. Case 2 is when encs prefers gzip encoding,
// the response is greater than 512 bytes and the response's content
// type is in ctMap.
//
// A gzipResponseWriter sets the Content-Encoding and Content-Type
// headers when appropriate. It is important to call the Close method
// when writing is finished in order to flush and close the
// gzipResponseWriter. The slice encs must contain only encodings from
// {encGzip,encIdentity} and contain at least one encoding.
type gzipResponseWriter struct {
	http.ResponseWriter
	httpStatus int
	ctMap      map[string]struct{}
	encs       []encoding
	gw         *gzip.Writer
	buf        *bytes.Buffer
}

func newGzipResponseWriter(w http.ResponseWriter, ctMap map[string]struct{}, encs []encoding) *gzipResponseWriter {
	buf := gzipBufPool.Get().(*bytes.Buffer)
	buf.Reset()
	return &gzipResponseWriter{
		ResponseWriter: w,
		httpStatus:     http.StatusOK,
		ctMap:          ctMap,
		encs:           encs,
		buf:            buf}
}

// init gets called by Write once at least 512 bytes have been written
// to the temporary buffer buf, or by Close if it has not yet been
// called. Firstly it determines the content type, either from the
// Content-Type header, or by calling http.DetectContentType on
// buf. Then, if needed, a gzip.Writer is initialized. Lastly,
// appropriate headers are set and the ResponseWriter's WriteHeader
// method is called.
func (w *gzipResponseWriter) init() {
	cth := w.Header().Get("Content-Type")
	var ct string
	if cth != "" {
		ct = cth
	} else {
		ct = http.DetectContentType(w.buf.Bytes())
	}
	var gzipContentType bool
	if mt, _, err := mime.ParseMediaType(ct); err == nil {
		if _, ok := w.ctMap[mt]; ok {
			gzipContentType = true
		}
	}
	var useGzip bool
	if w.Header().Get("Content-Encoding") == "" && w.encs[0] == encGzip {
		if gzipContentType && w.buf.Len() >= 512 || len(w.encs) == 1 {
			useGzip = true
		}
	}
	if useGzip {
		w.gw = gzipWriterPool.Get().(*gzip.Writer)
		w.gw.Reset(w.ResponseWriter)
		w.Header().Del("Content-Length")
		w.Header().Set("Content-Encoding", "gzip")
	}
	w.Header().Del("Accept-Ranges")
	if cth == "" {
		w.Header().Set("Content-Type", ct)
	}
	w.ResponseWriter.WriteHeader(w.httpStatus)
}

func (w *gzipResponseWriter) Write(p []byte) (int, error) {
	var n, written int
	var err error
	if w.buf != nil {
		written = w.buf.Len()
		_, _ = w.buf.Write(p)
		if w.buf.Len() < 512 {
			return len(p), nil
		}
		w.init()
		p = w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
	}
	switch {
	case w.gw != nil:
		n, err = w.gw.Write(p)
	default:
		n, err = w.ResponseWriter.Write(p)
	}
	n -= written
	if n < 0 {
		n = 0
	}
	return n, err
}

func (w *gzipResponseWriter) WriteHeader(httpStatus int) {
	// postpone WriteHeader call until end of init method
	w.httpStatus = httpStatus
}

func (w *gzipResponseWriter) Close() (err error) {
	if w.buf != nil {
		w.init()
		p := w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
		switch {
		case w.gw != nil:
			_, err = w.gw.Write(p)
		default:
			_, err = w.ResponseWriter.Write(p)
		}
	}
	if w.gw != nil {
		e := w.gw.Close()
		if e != nil {
			err = e
		}
		gzipWriterPool.Put(w.gw)
		w.gw = nil
	}
	return err
}

// An encoding is a supported content coding.
type encoding int

const (
	encIdentity encoding = iota
	encGzip
)

// acceptedEncodings returns the supported content codings that are
// accepted by the request r. It returns a slice of encodings in
// client preference order.
//
// If the Sec-WebSocket-Key header is present then compressed content
// encodings are not considered.
//
// ref: http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html
func acceptedEncodings(r *http.Request) []encoding {
	h := r.Header.Get("Accept-Encoding")
	swk := r.Header.Get("Sec-WebSocket-Key")
	if h == "" {
		return []encoding{encIdentity}
	}
	gzip := float64(-1)    // -1 means not accepted, 0 -> 1 means value of q
	identity := float64(0) // -1 means not accepted, 0 -> 1 means value of q
	for _, s := range strings.Split(h, ",") {
		f := strings.Split(s, ";")
		f0 := strings.ToLower(strings.Trim(f[0], " "))
		q := float64(1.0)
		if len(f) > 1 {
			f1 := strings.ToLower(strings.Trim(f[1], " "))
			if strings.HasPrefix(f1, "q=") {
				if flt, err := strconv.ParseFloat(f1[2:], 32); err == nil {
					if flt >= 0 && flt <= 1 {
						q = flt
					}
				}
			}
		}
		if (f0 == "gzip" || f0 == "*") && q > gzip && swk == "" {
			gzip = q
		}
		if (f0 == "gzip" || f0 == "*") && q == 0 {
			gzip = -1
		}
		if (f0 == "identity" || f0 == "*") && q > identity {
			identity = q
		}
		if (f0 == "identity" || f0 == "*") && q == 0 {
			identity = -1
		}
	}
	switch {
	case gzip == -1 && identity == -1:
		return []encoding{}
	case gzip == -1:
		return []encoding{encIdentity}
	case identity == -1:
		return []encoding{encGzip}
	case identity > gzip:
		return []encoding{encIdentity, encGzip}
	default:
		return []encoding{encGzip, encIdentity}
	}
}

// NewHandler returns a new http.Handler which wraps a handler h
// adding gzip compression to certain responses. There are two cases
// where gzip compression is done. Case 1 is responses whose requests
// only allow gzip encoding and forbid identity encoding (identity
// encoding meaning no encoding). Case 2 is responses whose requests
// prefer gzip encoding, whose size is greater than 512 bytes and
// whose content types are in contentTypes. If contentTypes is nil
// then DefaultContentTypes is considered instead.
//
// The new http.Handler sets the Content-Encoding, Vary and
// Content-Type headers in its responses as appropriate. If a request
// expresses a preference for gzip encoding then any Range headers are
// removed from the request before it is passed through to h and
// Accept-Ranges headers are stripped from corresponding
// responses. This happens regardless of whether gzip encoding is
// eventually used in the response or not.
func NewHandler(h http.Handler, contentTypes []string) http.Handler {
	if contentTypes == nil {
		contentTypes = DefaultContentTypes
	}
	ctMap := map[string]struct{}{}
	for _, ct := range contentTypes {
		ctMap[ct] = struct{}{}
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// add Vary header
		w.Header().Add("Vary", "Accept-Encoding")
		// check client's accepted encodings
		encs := acceptedEncodings(r)
		// return if no acceptable encodings
		if len(encs) == 0 {
			w.WriteHeader(http.StatusNotAcceptable)
			return
		}
		if encs[0] == encGzip {
			// cannot accept Range requests for possibly gzipped
			// responses
			r.Header.Del("Range")
			// create new ResponseWriter
			w = newGzipResponseWriter(w, ctMap, encs)
			defer w.(*gzipResponseWriter).Close()
		}
		// call original handler's ServeHTTP
		h.ServeHTTP(w, r)
	})
}
<commit_msg>Specify exactly when Content-Type headers are added<commit_after>/*
   Copyright 2015 The Httpgzip Authors. See the AUTHORS file at the
   top-level directory of this distribution and at
   <https://xi2.org/x/httpgzip/m/AUTHORS>.

   This file is part of Httpgzip.

   Httpgzip is free software: you can redistribute it and/or modify it
   under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   Httpgzip is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
   General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with Httpgzip.  If not, see <https://www.gnu.org/licenses/>.
*/

// Package httpgzip implements an http.Handler wrapper adding gzip
// compression for appropriate requests.
//
// It attempts to properly parse the request's Accept-Encoding header
// according to RFC 2616 and does not do a simple string search for
// "gzip" (which will fail to do the correct thing for values such as
// "*" or "identity,gzip;q=0"). It will serve either gzip or identity
// content codings, or return 406 Not Acceptable status if it can do
// neither.
//
// It works correctly with handlers which honour Range request headers
// (such as http.FileServer) by removing the Range header for requests
// which prefer gzip encoding. This is necessary since Range requests
// apply to the gzipped content but the wrapped handler is not aware
// of the compression when it writes byte ranges. The Accept-Ranges
// header is also stripped from corresponding responses.
//
// For requests which prefer gzip encoding a Content-Type header is
// set using http.DetectContentType if it is not set by the wrapped
// handler.
package httpgzip // import "xi2.org/x/httpgzip"

import (
	"bytes"
	"compress/gzip"
	"mime"
	"net/http"
	"strconv"
	"strings"
	"sync"
)

// DefaultContentTypes is the default list of content types for which
// a Handler considers gzip compression. This list originates from the
// file compression.conf within the Apache configuration found at
// https://html5boilerplate.com/.
var DefaultContentTypes = []string{
	"application/atom+xml",
	"application/javascript",
	"application/json",
	"application/ld+json",
	"application/manifest+json",
	"application/rdf+xml",
	"application/rss+xml",
	"application/schema+json",
	"application/vnd.geo+json",
	"application/vnd.ms-fontobject",
	"application/x-font-ttf",
	"application/x-javascript",
	"application/x-web-app-manifest+json",
	"application/xhtml+xml",
	"application/xml",
	"font/eot",
	"font/opentype",
	"image/bmp",
	"image/svg+xml",
	"image/vnd.microsoft.icon",
	"image/x-icon",
	"text/cache-manifest",
	"text/css",
	"text/html",
	"text/javascript",
	"text/plain",
	"text/vcard",
	"text/vnd.rim.location.xloc",
	"text/vtt",
	"text/x-component",
	"text/x-cross-domain-policy",
	"text/xml",
}

var gzipWriterPool = sync.Pool{
	New: func() interface{} { return gzip.NewWriter(nil) },
}

var gzipBufPool = sync.Pool{
	New: func() interface{} { return new(bytes.Buffer) },
}

// A gzipResponseWriter is a modified http.ResponseWriter. It adds
// gzip compression to certain responses, and there are two cases
// where this is done. Case 1 is when encs only allows gzip encoding
// and forbids identity. Case 2 is when encs prefers gzip encoding,
// the response is greater than 512 bytes and the response's content
// type is in ctMap.
//
// A gzipResponseWriter sets the Content-Encoding and Content-Type
// headers when appropriate. It is important to call the Close method
// when writing is finished in order to flush and close the
// gzipResponseWriter. The slice encs must contain only encodings from
// {encGzip,encIdentity} and contain at least one encoding.
type gzipResponseWriter struct {
	http.ResponseWriter
	httpStatus int
	ctMap      map[string]struct{}
	encs       []encoding
	gw         *gzip.Writer
	buf        *bytes.Buffer
}

func newGzipResponseWriter(w http.ResponseWriter, ctMap map[string]struct{}, encs []encoding) *gzipResponseWriter {
	buf := gzipBufPool.Get().(*bytes.Buffer)
	buf.Reset()
	return &gzipResponseWriter{
		ResponseWriter: w,
		httpStatus:     http.StatusOK,
		ctMap:          ctMap,
		encs:           encs,
		buf:            buf}
}

// init gets called by Write once at least 512 bytes have been written
// to the temporary buffer buf, or by Close if it has not yet been
// called. Firstly it determines the content type, either from the
// Content-Type header, or by calling http.DetectContentType on
// buf. Then, if needed, a gzip.Writer is initialized. Lastly,
// appropriate headers are set and the ResponseWriter's WriteHeader
// method is called.
func (w *gzipResponseWriter) init() {
	cth := w.Header().Get("Content-Type")
	var ct string
	if cth != "" {
		ct = cth
	} else {
		ct = http.DetectContentType(w.buf.Bytes())
	}
	var gzipContentType bool
	if mt, _, err := mime.ParseMediaType(ct); err == nil {
		if _, ok := w.ctMap[mt]; ok {
			gzipContentType = true
		}
	}
	var useGzip bool
	if w.Header().Get("Content-Encoding") == "" && w.encs[0] == encGzip {
		if gzipContentType && w.buf.Len() >= 512 || len(w.encs) == 1 {
			useGzip = true
		}
	}
	if useGzip {
		w.gw = gzipWriterPool.Get().(*gzip.Writer)
		w.gw.Reset(w.ResponseWriter)
		w.Header().Del("Content-Length")
		w.Header().Set("Content-Encoding", "gzip")
	}
	w.Header().Del("Accept-Ranges")
	if cth == "" {
		w.Header().Set("Content-Type", ct)
	}
	w.ResponseWriter.WriteHeader(w.httpStatus)
}

func (w *gzipResponseWriter) Write(p []byte) (int, error) {
	var n, written int
	var err error
	if w.buf != nil {
		written = w.buf.Len()
		_, _ = w.buf.Write(p)
		if w.buf.Len() < 512 {
			return len(p), nil
		}
		w.init()
		p = w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
	}
	switch {
	case w.gw != nil:
		n, err = w.gw.Write(p)
	default:
		n, err = w.ResponseWriter.Write(p)
	}
	n -= written
	if n < 0 {
		n = 0
	}
	return n, err
}

func (w *gzipResponseWriter) WriteHeader(httpStatus int) {
	// postpone WriteHeader call until end of init method
	w.httpStatus = httpStatus
}

func (w *gzipResponseWriter) Close() (err error) {
	if w.buf != nil {
		w.init()
		p := w.buf.Bytes()
		defer func() {
			gzipBufPool.Put(w.buf)
			w.buf = nil
		}()
		switch {
		case w.gw != nil:
			_, err = w.gw.Write(p)
		default:
			_, err = w.ResponseWriter.Write(p)
		}
	}
	if w.gw != nil {
		e := w.gw.Close()
		if e != nil {
			err = e
		}
		gzipWriterPool.Put(w.gw)
		w.gw = nil
	}
	return err
}

// An encoding is a supported content coding.
type encoding int

const (
	encIdentity encoding = iota
	encGzip
)

// acceptedEncodings returns the supported content codings that are
// accepted by the request r. It returns a slice of encodings in
// client preference order.
//
// If the Sec-WebSocket-Key header is present then compressed content
// encodings are not considered.
//
// ref: http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html
func acceptedEncodings(r *http.Request) []encoding {
	h := r.Header.Get("Accept-Encoding")
	swk := r.Header.Get("Sec-WebSocket-Key")
	if h == "" {
		return []encoding{encIdentity}
	}
	gzip := float64(-1)    // -1 means not accepted, 0 -> 1 means value of q
	identity := float64(0) // -1 means not accepted, 0 -> 1 means value of q
	for _, s := range strings.Split(h, ",") {
		f := strings.Split(s, ";")
		f0 := strings.ToLower(strings.Trim(f[0], " "))
		q := float64(1.0)
		if len(f) > 1 {
			f1 := strings.ToLower(strings.Trim(f[1], " "))
			if strings.HasPrefix(f1, "q=") {
				if flt, err := strconv.ParseFloat(f1[2:], 32); err == nil {
					if flt >= 0 && flt <= 1 {
						q = flt
					}
				}
			}
		}
		if (f0 == "gzip" || f0 == "*") && q > gzip && swk == "" {
			gzip = q
		}
		if (f0 == "gzip" || f0 == "*") && q == 0 {
			gzip = -1
		}
		if (f0 == "identity" || f0 == "*") && q > identity {
			identity = q
		}
		if (f0 == "identity" || f0 == "*") && q == 0 {
			identity = -1
		}
	}
	switch {
	case gzip == -1 && identity == -1:
		return []encoding{}
	case gzip == -1:
		return []encoding{encIdentity}
	case identity == -1:
		return []encoding{encGzip}
	case identity > gzip:
		return []encoding{encIdentity, encGzip}
	default:
		return []encoding{encGzip, encIdentity}
	}
}

// NewHandler returns a new http.Handler which wraps a handler h
// adding gzip compression to certain responses. There are two cases
// where gzip compression is done. Case 1 is responses whose requests
// only allow gzip encoding and forbid identity encoding (identity
// encoding meaning no encoding). Case 2 is responses whose requests
// prefer gzip encoding, whose size is greater than 512 bytes and
// whose content types are in contentTypes. If contentTypes is nil
// then DefaultContentTypes is considered instead.
//
// The new http.Handler sets the Content-Encoding, Vary and
// Content-Type headers in its responses as appropriate. If a request
// expresses a preference for gzip encoding then any Range headers are
// removed from the request before it is passed through to h and
// Accept-Ranges headers are stripped from corresponding
// responses. This happens regardless of whether gzip encoding is
// eventually used in the response or not.
func NewHandler(h http.Handler, contentTypes []string) http.Handler {
	if contentTypes == nil {
		contentTypes = DefaultContentTypes
	}
	ctMap := map[string]struct{}{}
	for _, ct := range contentTypes {
		ctMap[ct] = struct{}{}
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// add Vary header
		w.Header().Add("Vary", "Accept-Encoding")
		// check client's accepted encodings
		encs := acceptedEncodings(r)
		// return if no acceptable encodings
		if len(encs) == 0 {
			w.WriteHeader(http.StatusNotAcceptable)
			return
		}
		if encs[0] == encGzip {
			// cannot accept Range requests for possibly gzipped
			// responses
			r.Header.Del("Range")
			// create new ResponseWriter
			w = newGzipResponseWriter(w, ctMap, encs)
			defer w.(*gzipResponseWriter).Close()
		}
		// call original handler's ServeHTTP
		h.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>package authn

import (
	"encoding/json"
	"io/ioutil"
	"sync"
	"time"
	"util"

	fernet "github.com/fernet/fernet-go"
)

var mint struct {
	*fernet.Key
	*sync.RWMutex
}

func InitMint() {
	confVersion = util.GetConfig().Version
	loadMintKey()
}

func loadMintKey() {
	LoadMintKeyByName(util.GetConfig().Authn.MintKeyName)
	confVersion = util.GetConfig().Version
}

// loadMintKeyByName loads a key by filename and strores it in the struct
// The function is threadsafe and panics if the key file is invalid
func LoadMintKeyByName(filename string) {

	keyPath := util.GetPaths([]string{filename})[0]

	log.Debugf("Attempting to load mint key from %s", keyPath)
	b, err := ioutil.ReadFile(keyPath)

	if err != nil {
		panic(err)
	}

	k, err := fernet.DecodeKey(string(b))

	if err != nil {
		panic(err)
	}

	log.Debugf("Successfully loaded mint key from %s", keyPath)
	// Store only after we are sure loading was good
	mint.Lock()
	defer mint.Unlock()
	mint.Key = k
}

// GetToken wraps the incoming username into a TokenStruct, serializes the result to json
// and generates a Fernet token based on the resulting string
func GetToken(username string) string {
	var t []byte

	// If the configuration has changed, re-load the keys
	if confVersion != util.GetConfig().Version {
		loadMintKey()
	}

	n, _ := time.Now().MarshalText()
	t, _ = json.Marshal(TokenStruct{Username: username, MintTime: string(n)})

	token, err := fernet.EncryptAndSign(t, GetKey())
	if err != nil {
		panic(err)
	}
	return string(token)
}

// GetKey returns the current key used for session tokens
// If the key not initialized, nil is returned
func GetKey() *fernet.Key {
	mint.RLock()
	defer mint.RUnlock()
	return mint.Key

}

// ReissueToken re-issues a token based on a previous valid token
func ReissueToken(token string) (string, error) {
	var newToken string
	var decryptedTokensUsername string
	var e error

	if decryptedTokensUsername, e = Validate(token); e != nil {
		log.Info("Reissue request for invalid token " + token)
		return "", e
	}

	newToken = GetToken(decryptedTokensUsername)

	return newToken, nil
}
<commit_msg>fix anonymous struct problem in Mint so tests succeed<commit_after>package authn

import (
	"encoding/json"
	"io/ioutil"
	"sync"
	"time"
	"util"

	fernet "github.com/fernet/fernet-go"
)

var mint struct {
	*fernet.Key
	sync.RWMutex
}

func InitMint() {
	confVersion = util.GetConfig().Version
	loadMintKey()
}

func loadMintKey() {
	LoadMintKeyByName(util.GetConfig().Authn.MintKeyName)
	confVersion = util.GetConfig().Version
}

// loadMintKeyByName loads a key by filename and strores it in the struct
// The function is threadsafe and panics if the key file is invalid
func LoadMintKeyByName(filename string) {

	keyPath := util.GetPaths([]string{filename})[0]

	log.Debugf("Attempting to load mint key from %s", keyPath)
	b, err := ioutil.ReadFile(keyPath)

	if err != nil {
		panic(err)
	}

	k, err := fernet.DecodeKey(string(b))

	if err != nil {
		panic(err)
	}

	log.Debugf("Successfully loaded mint key from %s", keyPath)
	// Store only after we are sure loading was good
	mint.Lock()
	defer mint.Unlock()
	mint.Key = k
}

// GetToken wraps the incoming username into a TokenStruct, serializes the result to json
// and generates a Fernet token based on the resulting string
func GetToken(username string) string {
	var t []byte

	// If the configuration has changed, re-load the keys
	if confVersion != util.GetConfig().Version {
		loadMintKey()
	}

	n, _ := time.Now().MarshalText()
	t, _ = json.Marshal(TokenStruct{Username: username, MintTime: string(n)})

	token, err := fernet.EncryptAndSign(t, GetKey())
	if err != nil {
		panic(err)
	}
	return string(token)
}

// GetKey returns the current key used for session tokens
// If the key not initialized, nil is returned
func GetKey() *fernet.Key {
	mint.RLock()
	defer mint.RUnlock()
	return mint.Key

}

// ReissueToken re-issues a token based on a previous valid token
func ReissueToken(token string) (string, error) {
	var newToken string
	var decryptedTokensUsername string
	var e error

	if decryptedTokensUsername, e = Validate(token); e != nil {
		log.Info("Reissue request for invalid token " + token)
		return "", e
	}

	newToken = GetToken(decryptedTokensUsername)

	return newToken, nil
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
)

var cmdRun = &Command{
	UsageLine: "run [build flags] gofiles... [arguments...]",
	Short:     "compile and run Go program",
	Long: `
Run compiles and runs the main package comprising the named Go source files.

For more about build flags, see 'go help build'.

See also: go build.
	`,
}

func init() {
	cmdRun.Run = runRun // break init loop

	addBuildFlags(cmdRun)
}

func printStderr(args ...interface{}) (int, error) {
	return fmt.Fprint(os.Stderr, args...)
}

func runRun(cmd *Command, args []string) {
	var b builder
	b.init()
	b.print = printStderr
	i := 0
	for i < len(args) && strings.HasSuffix(args[i], ".go") {
		i++
	}
	files, cmdArgs := args[:i], args[i:]
	p := goFilesPackage(files)
	if p.Error != nil {
		fatalf("%s", p.Error)
	}
	if p.Name != "main" {
		fatalf("cannot run non-main package")
	}
	p.target = "" // must build - not up to date
	a1 := b.action(modeBuild, modeBuild, p)
	a := &action{f: (*builder).runProgram, args: cmdArgs, deps: []*action{a1}}
	b.do(a)
}

// runProgram is the action for running a binary that has already
// been compiled.  We ignore exit status.
func (b *builder) runProgram(a *action) error {
	if buildN || buildX {
		b.showcmd("", "%s %s", a.deps[0].target, strings.Join(a.args, " "))
		if buildN {
			return nil
		}
	}

	runStdin(a.deps[0].target, a.args)
	return nil
}

// runStdin is like run, but connects Stdin.
func runStdin(cmdargs ...interface{}) {
	cmdline := stringList(cmdargs...)
	cmd := exec.Command(cmdline[0], cmdline[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		errorf("%v", err)
	}
}
<commit_msg>cmd/go: fix run errors<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
)

var cmdRun = &Command{
	UsageLine: "run [build flags] gofiles... [arguments...]",
	Short:     "compile and run Go program",
	Long: `
Run compiles and runs the main package comprising the named Go source files.

For more about build flags, see 'go help build'.

See also: go build.
	`,
}

func init() {
	cmdRun.Run = runRun // break init loop

	addBuildFlags(cmdRun)
}

func printStderr(args ...interface{}) (int, error) {
	return fmt.Fprint(os.Stderr, args...)
}

func runRun(cmd *Command, args []string) {
	var b builder
	b.init()
	b.print = printStderr
	i := 0
	for i < len(args) && strings.HasSuffix(args[i], ".go") {
		i++
	}
	files, cmdArgs := args[:i], args[i:]
	if len(files) == 0 {
		fatalf("go run: no go files listed")
	}
	p := goFilesPackage(files)
	if p.Error != nil {
		fatalf("%s", p.Error)
	}
	if p.Name != "main" {
		fatalf("go run: cannot run non-main package")
	}
	p.target = "" // must build - not up to date
	a1 := b.action(modeBuild, modeBuild, p)
	a := &action{f: (*builder).runProgram, args: cmdArgs, deps: []*action{a1}}
	b.do(a)
}

// runProgram is the action for running a binary that has already
// been compiled.  We ignore exit status.
func (b *builder) runProgram(a *action) error {
	if buildN || buildX {
		b.showcmd("", "%s %s", a.deps[0].target, strings.Join(a.args, " "))
		if buildN {
			return nil
		}
	}

	runStdin(a.deps[0].target, a.args)
	return nil
}

// runStdin is like run, but connects Stdin.
func runStdin(cmdargs ...interface{}) {
	cmdline := stringList(cmdargs...)
	cmd := exec.Command(cmdline[0], cmdline[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		errorf("%v", err)
	}
}
<|endoftext|>