<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"log"
	"net/url"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/valyala/bytebufferpool"
	"github.com/valyala/fasthttp"
	"github.com/xtrafrancyz/vk-proxy/bytefmt"
	"github.com/xtrafrancyz/vk-proxy/replacer"
)

const (
	readBufferSize = 8192
)

var (
	gzip            = []byte("gzip")
	vkProxyName     = []byte("vk-proxy")
	serverHeader    = []byte("Server")
	setCookie       = []byte("Set-Cookie")
	acceptEncoding  = []byte("Accept-Encoding")
	contentEncoding = []byte("Content-Encoding")
)

type ProxyConfig struct {
	ReduceMemoryUsage bool
	BaseDomain        string
	BaseStaticDomain  string
	LogVerbosity      int
	GzipUpstream      bool
	FilterFeed        bool
}

type Proxy struct {
	client   *fasthttp.Client
	server   *fasthttp.Server
	replacer *replacer.Replacer
	tracker  *tracker
	config   ProxyConfig
}

func NewProxy(config ProxyConfig) *Proxy {
	p := &Proxy{
		client: &fasthttp.Client{
			Name:           "vk-proxy",
			ReadBufferSize: readBufferSize,
			TLSConfig:      &tls.Config{InsecureSkipVerify: true},
		},
		replacer: &replacer.Replacer{
			ProxyBaseDomain:   config.BaseDomain,
			ProxyStaticDomain: config.BaseStaticDomain,
		},
		tracker: &tracker{
			uniqueUsers: make(map[string]bool),
		},
		config: config,
	}
	p.server = &fasthttp.Server{
		Handler:           p.handleProxy,
		ReduceMemoryUsage: config.ReduceMemoryUsage,
		ReadBufferSize:    readBufferSize,
		Name:              "vk-proxy",
	}
	if p.config.LogVerbosity > 0 {
		p.tracker.start()
	}
	return p
}

func (p *Proxy) ListenTCP(host string) error {
	log.Printf("Starting server on http://%s", host)
	return p.server.ListenAndServe(host)
}

func (p *Proxy) ListenUnix(path string) error {
	log.Printf("Starting server on http://unix:%s", path)
	return p.server.ListenAndServeUNIX(path, 0777)
}

func (p *Proxy) handleProxy(ctx *fasthttp.RequestCtx) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic when proxying the request: %s%s", r, debug.Stack())
			ctx.Error("500 Internal Server Error", 500)
		}
	}()
	start := time.Now()

	replaceContext := &replacer.ReplaceContext{
		Method:     ctx.Method(),
		FilterFeed: p.config.FilterFeed,
	}

	if !p.prepareProxyRequest(ctx, replaceContext) {
		ctx.Error("400 Bad Request", 400)
		return
	}

	if replaceContext.Host == "api.vk.com" &&
		(replaceContext.Path == "/away" || replaceContext.Path == "/away.php") {
		p.handleAway(ctx)
		return
	}

	err := p.client.DoTimeout(&ctx.Request, &ctx.Response, 30*time.Second)
	if err == nil {
		err = p.processProxyResponse(ctx, replaceContext)
	}

	elapsed := time.Since(start).Round(100 * time.Microsecond)

	if err != nil {
		log.Printf("%s %s error: %s", elapsed, ctx.Path(), err)
		if strings.Contains(err.Error(), "timed out") || strings.Contains(err.Error(), "timeout") {
			ctx.Error("408 Request Timeout", 408)
		} else {
			ctx.Error("500 Internal Server Error", 500)
		}
		return
	}

	if p.config.LogVerbosity > 0 {
		ip := ctx.Request.Header.Peek("CF-Connecting-IP") // Cloudflare
		if ip == nil {
			ip = ctx.Request.Header.Peek("X-Real-IP") // nginx
		}
		if ip == nil {
			ip = []byte(ctx.RemoteIP().String()) // real
		}
		p.tracker.trackRequest(string(ip), len(ctx.Response.Body()))
	}

	if p.config.LogVerbosity == 2 {
		log.Printf("%s %s %s%s %s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))))
	} else if p.config.LogVerbosity == 3 {
		log.Printf("%s %s %s%s %s\n%s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))), ctx.Response.Body())
	}
}

func (p *Proxy) handleAway(ctx *fasthttp.RequestCtx) {
	to := string(ctx.QueryArgs().Peek("to"))
	if to == "" {
		ctx.Error("Bad Request: 'to' argument is not set", 400)
		return
	}
	to, err := url.QueryUnescape(to)
	if err != nil {
		ctx.Error("Bad Request: could not unescape url", 400)
		return
	}
	ctx.Redirect(to, fasthttp.StatusMovedPermanently)
}

func (p *Proxy) prepareProxyRequest(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) bool {
	// Routing
	req := &ctx.Request
	uri := string(req.RequestURI())
	host := ""
	if strings.HasPrefix(uri, "/@") {
		slashIndex := strings.IndexByte(uri[2:], '/')
		if slashIndex == -1 {
			return false
		}
		endpoint := uri[2 : slashIndex+2]
		if endpoint != "vk.com" &&
			!strings.HasSuffix(endpoint, ".vk.com") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.net") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.com") {
			return false
		}
		host = endpoint
		uri = uri[2+slashIndex:]
		req.SetRequestURI(uri)
	} else if altHost := req.Header.Peek("Proxy-Host"); altHost != nil {
		host = string(altHost)
		switch host {
		case "static.vk.com":
		case "oauth.vk.com":
		default:
			return false
		}
		req.Header.Del("Proxy-Host")
	} else {
		host = "api.vk.com"
	}
	req.SetHost(host)

	// Replace some request data
	replaceContext.Host = host
	replaceContext.Path = string(ctx.Path())
	p.replacer.DoReplaceRequest(req, replaceContext)

	// After req.URI() call it is impossible to modify URI
	req.URI().SetScheme("https")
	if p.config.GzipUpstream {
		req.Header.SetBytesKV(acceptEncoding, gzip)
	} else {
		req.Header.DelBytes(acceptEncoding)
	}
	return true
}

func (p *Proxy) processProxyResponse(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) error {
	res := &ctx.Response
	res.Header.DelBytes(setCookie)
	res.Header.SetBytesKV(serverHeader, vkProxyName)

	var buf *bytebufferpool.ByteBuffer
	// Gunzip body if needed
	if bytes.Contains(res.Header.PeekBytes(contentEncoding), gzip) {
		res.Header.DelBytes(contentEncoding)
		var err error
		buf = &bytebufferpool.ByteBuffer{}
		buf.B, err = res.BodyGunzip()
		if err != nil {
			return err
		}
	} else {
		buf = &bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		}
	}

	buf = p.replacer.DoReplaceResponse(res, buf, replaceContext)

	// avoid copying and save old buffer
	buf.B = res.SwapBody(buf.B)
	if cap(buf.B) > 10 {
		replacer.ReleaseBuffer(buf)
	}
	return nil
}

type tracker struct {
	lock        sync.Mutex
	requests    uint32
	bytes       uint64
	uniqueUsers map[string]bool
}

func (t *tracker) start() {
	go func() {
		for range time.Tick(60 * time.Second) {
			t.lock.Lock()
			log.Printf("Requests: %d, Traffic: %s, Online: %d", t.requests, bytefmt.ByteSize(t.bytes), len(t.uniqueUsers))
			t.requests = 0
			t.bytes = 0
			t.uniqueUsers = make(map[string]bool)
			t.lock.Unlock()
		}
	}()
}

func (t *tracker) trackRequest(ip string, size int) {
	t.lock.Lock()

	t.uniqueUsers[ip] = true
	t.requests++
	t.bytes += uint64(size)

	t.lock.Unlock()
}
<commit_msg>Оптимизирована аллокация памяти<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"log"
	"net/url"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/valyala/bytebufferpool"
	"github.com/valyala/fasthttp"
	"github.com/xtrafrancyz/vk-proxy/bytefmt"
	"github.com/xtrafrancyz/vk-proxy/replacer"
)

const (
	readBufferSize = 8192
)

var (
	gzip            = []byte("gzip")
	vkProxyName     = []byte("vk-proxy")
	serverHeader    = []byte("Server")
	setCookie       = []byte("Set-Cookie")
	acceptEncoding  = []byte("Accept-Encoding")
	contentEncoding = []byte("Content-Encoding")
)

type ProxyConfig struct {
	ReduceMemoryUsage bool
	BaseDomain        string
	BaseStaticDomain  string
	LogVerbosity      int
	GzipUpstream      bool
	FilterFeed        bool
}

type Proxy struct {
	client   *fasthttp.Client
	server   *fasthttp.Server
	replacer *replacer.Replacer
	tracker  *tracker
	config   ProxyConfig
}

func NewProxy(config ProxyConfig) *Proxy {
	p := &Proxy{
		client: &fasthttp.Client{
			Name:           "vk-proxy",
			ReadBufferSize: readBufferSize,
			TLSConfig:      &tls.Config{InsecureSkipVerify: true},
		},
		replacer: &replacer.Replacer{
			ProxyBaseDomain:   config.BaseDomain,
			ProxyStaticDomain: config.BaseStaticDomain,
		},
		tracker: &tracker{
			uniqueUsers: make(map[string]bool),
		},
		config: config,
	}
	p.server = &fasthttp.Server{
		Handler:           p.handleProxy,
		ReduceMemoryUsage: config.ReduceMemoryUsage,
		ReadBufferSize:    readBufferSize,
		Name:              "vk-proxy",
	}
	if p.config.LogVerbosity > 0 {
		p.tracker.start()
	}
	return p
}

func (p *Proxy) ListenTCP(host string) error {
	log.Printf("Starting server on http://%s", host)
	return p.server.ListenAndServe(host)
}

func (p *Proxy) ListenUnix(path string) error {
	log.Printf("Starting server on http://unix:%s", path)
	return p.server.ListenAndServeUNIX(path, 0777)
}

func (p *Proxy) handleProxy(ctx *fasthttp.RequestCtx) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic when proxying the request: %s%s", r, debug.Stack())
			ctx.Error("500 Internal Server Error", 500)
		}
	}()
	start := time.Now()

	replaceContext := &replacer.ReplaceContext{
		Method:     ctx.Method(),
		FilterFeed: p.config.FilterFeed,
	}

	if !p.prepareProxyRequest(ctx, replaceContext) {
		ctx.Error("400 Bad Request", 400)
		return
	}

	if replaceContext.Host == "api.vk.com" &&
		(replaceContext.Path == "/away" || replaceContext.Path == "/away.php") {
		p.handleAway(ctx)
		return
	}

	err := p.client.DoTimeout(&ctx.Request, &ctx.Response, 30*time.Second)
	if err == nil {
		err = p.processProxyResponse(ctx, replaceContext)
	}

	elapsed := time.Since(start).Round(100 * time.Microsecond)

	if err != nil {
		log.Printf("%s %s error: %s", elapsed, ctx.Path(), err)
		if strings.Contains(err.Error(), "timed out") || strings.Contains(err.Error(), "timeout") {
			ctx.Error("408 Request Timeout", 408)
		} else {
			ctx.Error("500 Internal Server Error", 500)
		}
		return
	}

	if p.config.LogVerbosity > 0 {
		ip := ctx.Request.Header.Peek("CF-Connecting-IP") // Cloudflare
		if ip == nil {
			ip = ctx.Request.Header.Peek("X-Real-IP") // nginx
		}
		if ip == nil {
			ip = []byte(ctx.RemoteIP().String()) // real
		}
		p.tracker.trackRequest(string(ip), len(ctx.Response.Body()))
	}

	if p.config.LogVerbosity == 2 {
		log.Printf("%s %s %s%s %s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))))
	} else if p.config.LogVerbosity == 3 {
		log.Printf("%s %s %s%s %s\n%s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))), ctx.Response.Body())
	}
}

func (p *Proxy) handleAway(ctx *fasthttp.RequestCtx) {
	to := string(ctx.QueryArgs().Peek("to"))
	if to == "" {
		ctx.Error("Bad Request: 'to' argument is not set", 400)
		return
	}
	to, err := url.QueryUnescape(to)
	if err != nil {
		ctx.Error("Bad Request: could not unescape url", 400)
		return
	}
	ctx.Redirect(to, fasthttp.StatusMovedPermanently)
}

func (p *Proxy) prepareProxyRequest(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) bool {
	// Routing
	req := &ctx.Request
	uri := string(req.RequestURI())
	host := ""
	if strings.HasPrefix(uri, "/@") {
		slashIndex := strings.IndexByte(uri[2:], '/')
		if slashIndex == -1 {
			return false
		}
		endpoint := uri[2 : slashIndex+2]
		if endpoint != "vk.com" &&
			!strings.HasSuffix(endpoint, ".vk.com") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.net") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.com") {
			return false
		}
		host = endpoint
		uri = uri[2+slashIndex:]
		req.SetRequestURI(uri)
	} else if altHost := req.Header.Peek("Proxy-Host"); altHost != nil {
		host = string(altHost)
		switch host {
		case "static.vk.com":
		case "oauth.vk.com":
		default:
			return false
		}
		req.Header.Del("Proxy-Host")
	} else {
		host = "api.vk.com"
	}
	req.SetHost(host)

	// Replace some request data
	replaceContext.Host = host
	replaceContext.Path = string(ctx.Path())
	p.replacer.DoReplaceRequest(req, replaceContext)

	// After req.URI() call it is impossible to modify URI
	req.URI().SetScheme("https")
	if p.config.GzipUpstream {
		req.Header.SetBytesKV(acceptEncoding, gzip)
	} else {
		req.Header.DelBytes(acceptEncoding)
	}
	return true
}

func (p *Proxy) processProxyResponse(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) error {
	res := &ctx.Response
	res.Header.DelBytes(setCookie)
	res.Header.SetBytesKV(serverHeader, vkProxyName)

	var buf *bytebufferpool.ByteBuffer
	// Gunzip body if needed
	if bytes.Contains(res.Header.PeekBytes(contentEncoding), gzip) {
		res.Header.DelBytes(contentEncoding)
		buf = replacer.AcquireBuffer()
		_, err := fasthttp.WriteGunzip(buf, res.Body())
		if err != nil {
			replacer.ReleaseBuffer(buf)
			return err
		}
		replacer.ReleaseBuffer(&bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		})
	} else {
		buf = &bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		}
	}

	buf = p.replacer.DoReplaceResponse(res, buf, replaceContext)

	// avoid copying and save old buffer
	buf.B = res.SwapBody(buf.B)
	if cap(buf.B) > 10 {
		replacer.ReleaseBuffer(buf)
	}
	return nil
}

type tracker struct {
	lock        sync.Mutex
	requests    uint32
	bytes       uint64
	uniqueUsers map[string]bool
}

func (t *tracker) start() {
	go func() {
		for range time.Tick(60 * time.Second) {
			t.lock.Lock()
			log.Printf("Requests: %d, Traffic: %s, Online: %d", t.requests, bytefmt.ByteSize(t.bytes), len(t.uniqueUsers))
			t.requests = 0
			t.bytes = 0
			t.uniqueUsers = make(map[string]bool)
			t.lock.Unlock()
		}
	}()
}

func (t *tracker) trackRequest(ip string, size int) {
	t.lock.Lock()

	t.uniqueUsers[ip] = true
	t.requests++
	t.bytes += uint64(size)

	t.lock.Unlock()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/01org/cc-oci-runtime/proxy/api"

	"github.com/golang/glog"
)

// Main struct holding the proxy state
type proxy struct {
	// Protect concurrent accesses from separate client goroutines to this
	// structure fields
	sync.Mutex

	// proxy socket
	listener net.Listener

	// vms are hashed by their containerID
	vms map[string]*vm

	// Output the VM console on stderr
	enableVMConsole bool

	wg sync.WaitGroup
}

// Represents a client, either a cc-oci-runtime or cc-shim process having
// opened a socket to the proxy
type client struct {
	id    uint64
	proxy *proxy
	vm    *vm

	conn net.Conn
}

func (c *client) info(lvl glog.Level, msg string) {
	if !glog.V(lvl) {
		return
	}
	glog.Infof("[client #%d] %s", c.id, msg)
}

func (c *client) infof(lvl glog.Level, fmt string, a ...interface{}) {
	if !glog.V(lvl) {
		return
	}
	a = append(a, 0)
	copy(a[1:], a[0:])
	a[0] = c.id
	glog.Infof("[client #%d] "+fmt, a...)
}

// "hello"
func helloHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hello := api.Hello{}

	if err := json.Unmarshal(data, &hello); err != nil {
		response.SetError(err)
		return
	}

	if hello.ContainerID == "" || hello.CtlSerial == "" || hello.IoSerial == "" {
		response.SetErrorMsg("malformed hello command")
	}

	proxy := client.proxy
	proxy.Lock()
	if _, ok := proxy.vms[hello.ContainerID]; ok {

		proxy.Unlock()
		response.SetErrorf("%s: container already registered",
			hello.ContainerID)
		return
	}

	client.infof(1, "hello(containerId=%s,ctlSerial=%s,ioSerial=%s,console=%s)", hello.ContainerID,
		hello.CtlSerial, hello.IoSerial, hello.Console)

	vm := newVM(hello.ContainerID, hello.CtlSerial, hello.IoSerial)
	proxy.vms[hello.ContainerID] = vm
	proxy.Unlock()

	if hello.Console != "" && proxy.enableVMConsole {
		vm.setConsole(hello.Console)
	}

	if err := vm.Connect(); err != nil {
		proxy.Lock()
		delete(proxy.vms, hello.ContainerID)
		proxy.Unlock()
		response.SetError(err)
		return
	}

	client.vm = vm

	// We start one goroutine per-VM to monitor the qemu process
	proxy.wg.Add(1)
	go func() {
		<-vm.OnVMLost()
		vm.Close()
		proxy.wg.Done()
	}()
}

// "attach"
func attachHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	proxy := client.proxy

	attach := api.Attach{}
	if err := json.Unmarshal(data, &attach); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[attach.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", attach.ContainerID)
		return
	}

	client.infof(1, "attach(containerId=%s)", attach.ContainerID)

	client.vm = vm
}

// "bye"
func byeHandler(data []byte, userData interface{}, response *handlerResponse) {
	// Bye only affects the proxy.vms map and so removes the VM from the
	// client visible API.
	// vm.Close(), which tears down the VM object, is done at the end of
	// the VM life cycle, when  we detect the qemu process is effectively
	// gone (see helloHandler)

	client := userData.(*client)
	proxy := client.proxy

	bye := api.Bye{}
	if err := json.Unmarshal(data, &bye); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[bye.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", bye.ContainerID)
		return
	}

	client.info(1, "bye()")

	proxy.Lock()
	delete(proxy.vms, vm.containerID)
	proxy.Unlock()

	client.vm = nil
}

// "allocateIO"
func allocateIoHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	vm := client.vm

	allocateIo := api.AllocateIo{}
	if err := json.Unmarshal(data, &allocateIo); err != nil {
		response.SetError(err)
		return
	}

	if allocateIo.NStreams < 1 || allocateIo.NStreams > 2 {
		response.SetErrorf("asking for unexpected number of streams (%d)",
			allocateIo.NStreams)
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "allocateIo(nStreams=%d)", allocateIo.NStreams)

	// We'll send c0 to the client, keep c1
	c0, c1, err := Socketpair()
	if err != nil {
		response.SetError(err)
		return
	}

	f0, err := c0.File()
	if err != nil {
		response.SetError(err)
		return
	}

	ioBase := vm.AllocateIo(allocateIo.NStreams, client.id, c1)

	client.infof(1, "-> %d streams allocated, ioBase=%d", allocateIo.NStreams, ioBase)

	response.AddResult("ioBase", ioBase)
	response.SetFile(f0)

	// File() dups the underlying fd, so it's safe to close c0 here (will
	// keep the c0 <-> c1 connection alive).
	c0.Close()
}

// "hyper"
func hyperHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hyper := api.Hyper{}
	vm := client.vm

	if err := json.Unmarshal(data, &hyper); err != nil {
		response.SetError(err)
		return
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "hyper(cmd=%s, data=%s)", hyper.HyperName, hyper.Data)

	err := vm.SendMessage(hyper.HyperName, hyper.Data)
	response.SetError(err)
}

func newProxy() *proxy {
	return &proxy{
		vms: make(map[string]*vm),
	}
}

// DefaultSocketPath is populated at link time with the value of:
//   ${locatestatedir}/run/cc-oci-runtime/proxy
var DefaultSocketPath string

// ArgSocketPath is populated at runtime from the option -socket-path
var ArgSocketPath = flag.String("socket-path", "", "specify path to socket file")

func (proxy *proxy) init() error {
	var l net.Listener
	var err error

	// flags
	v := flag.Lookup("v").Value.(flag.Getter).Get().(glog.Level)
	proxy.enableVMConsole = v >= 3

	// Open the proxy socket
	fds := listenFds()

	if len(fds) > 1 {
		return fmt.Errorf("too many activated sockets (%d)", len(fds))
	} else if len(fds) == 1 {
		fd := fds[0]
		l, err = net.FileListener(fd)
		if err != nil {
			return fmt.Errorf("couldn't listen on socket: %v", err)
		}
	} else {
		// Invoking "go build" without any linker option will not
		// populate DefaultSocketPath, so fallback to a reasonable
		// path.
		if DefaultSocketPath == "" {
			DefaultSocketPath = "/var/run/cc-oci-runtime/proxy.sock"
		}

		socketPath := DefaultSocketPath
		if len(*ArgSocketPath) != 0 {
			socketPath = *ArgSocketPath
		}

		socketDir := filepath.Dir(socketPath)
		if err = os.MkdirAll(socketDir, 0750); err != nil {
			return fmt.Errorf("couldn't create socket directory: %v", err)
		}
		if err = os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("couldn't remove exiting socket: %v", err)
		}
		l, err = net.ListenUnix("unix", &net.UnixAddr{Name: socketPath, Net: "unix"})
		if err != nil {
			return fmt.Errorf("couldn't create AF_UNIX socket: %v", err)
		}
		if err = os.Chmod(socketPath, 0660|os.ModeSocket); err != nil {
			return fmt.Errorf("couldn't set mode on socket: %v", err)
		}
	}

	proxy.listener = l

	return nil
}

var nextClientID = uint64(1)

func (proxy *proxy) serveNewClient(proto *protocol, newConn net.Conn) {
	newClient := &client{
		id:    nextClientID,
		proxy: proxy,
		conn:  newConn,
	}

	atomic.AddUint64(&nextClientID, 1)

	// Unfortunately it's hard to find out information on the peer
	// at the other end of a unix socket. We use a per-client ID to
	// identify connections.
	newClient.info(1, "client connected")

	if err := proto.Serve(newConn, newClient); err != nil && err != io.EOF {
		newClient.infof(1, "error serving client: %v", err)
	}

	newConn.Close()
	newClient.info(1, "connection closed")
}

func (proxy *proxy) serve() {

	// Define the client (runtime/shim) <-> proxy protocol
	proto := newProtocol()
	proto.Handle("hello", helloHandler)
	proto.Handle("attach", attachHandler)
	proto.Handle("bye", byeHandler)
	proto.Handle("allocateIO", allocateIoHandler)
	proto.Handle("hyper", hyperHandler)

	glog.V(1).Info("proxy started")

	for {
		conn, err := proxy.listener.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "couldn't accept connection:", err)
			continue
		}

		go proxy.serveNewClient(proto, conn)
	}
}

func proxyMain() {
	proxy := newProxy()
	if err := proxy.init(); err != nil {
		fmt.Fprintln(os.Stderr, "init:", err.Error())
		os.Exit(1)
	}
	proxy.serve()

	// Wait for all the goroutines started by helloHandler to finish.
	//
	// Not stricly necessary as:
	//   • currently proxy.serve() cannot return,
	//   • even if it was, the process is about to exit anyway...
	//
	// That said, this wait group is used in the tests to ensure proper
	// serialisation between runs of proxyMain()(see proxy/proxy_test.go).
	proxy.wg.Wait()
}

func initLogging() {
	// We print logs on stderr by default.
	flag.Set("logtostderr", "true")

	// It can be practical to use an environment variable to trigger a verbose output
	level := os.Getenv("CC_PROXY_LOG_LEVEL")
	if level != "" {
		flag.Set("v", level)
	}
}

type profiler struct {
	enabled bool
	host    string
	port    uint
}

func (p *profiler) setup() {
	if !p.enabled {
		return
	}

	addr := fmt.Sprintf("%s:%d", p.host, p.port)
	url := "http://" + addr + "/debug/pprof"
	glog.V(1).Info("pprof enabled on " + url)

	go func() {
		http.ListenAndServe(addr, nil)
	}()
}

func main() {
	var pprof profiler

	initLogging()

	flag.BoolVar(&pprof.enabled, "pprof", false,
		"enable pprof ")
	flag.StringVar(&pprof.host, "pprof-host", "localhost",
		"host the pprof server will be bound to")
	flag.UintVar(&pprof.port, "pprof-port", 6060,
		"port the pprof server will be bound to")

	flag.Parse()
	defer glog.Flush()

	pprof.setup()
	proxyMain()
}
<commit_msg>proxy: Add a debug message to print the proxy socket path<commit_after>// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/01org/cc-oci-runtime/proxy/api"

	"github.com/golang/glog"
)

// Main struct holding the proxy state
type proxy struct {
	// Protect concurrent accesses from separate client goroutines to this
	// structure fields
	sync.Mutex

	// proxy socket
	listener net.Listener

	// vms are hashed by their containerID
	vms map[string]*vm

	// Output the VM console on stderr
	enableVMConsole bool

	wg sync.WaitGroup
}

// Represents a client, either a cc-oci-runtime or cc-shim process having
// opened a socket to the proxy
type client struct {
	id    uint64
	proxy *proxy
	vm    *vm

	conn net.Conn
}

func (c *client) info(lvl glog.Level, msg string) {
	if !glog.V(lvl) {
		return
	}
	glog.Infof("[client #%d] %s", c.id, msg)
}

func (c *client) infof(lvl glog.Level, fmt string, a ...interface{}) {
	if !glog.V(lvl) {
		return
	}
	a = append(a, 0)
	copy(a[1:], a[0:])
	a[0] = c.id
	glog.Infof("[client #%d] "+fmt, a...)
}

// "hello"
func helloHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hello := api.Hello{}

	if err := json.Unmarshal(data, &hello); err != nil {
		response.SetError(err)
		return
	}

	if hello.ContainerID == "" || hello.CtlSerial == "" || hello.IoSerial == "" {
		response.SetErrorMsg("malformed hello command")
	}

	proxy := client.proxy
	proxy.Lock()
	if _, ok := proxy.vms[hello.ContainerID]; ok {

		proxy.Unlock()
		response.SetErrorf("%s: container already registered",
			hello.ContainerID)
		return
	}

	client.infof(1, "hello(containerId=%s,ctlSerial=%s,ioSerial=%s,console=%s)", hello.ContainerID,
		hello.CtlSerial, hello.IoSerial, hello.Console)

	vm := newVM(hello.ContainerID, hello.CtlSerial, hello.IoSerial)
	proxy.vms[hello.ContainerID] = vm
	proxy.Unlock()

	if hello.Console != "" && proxy.enableVMConsole {
		vm.setConsole(hello.Console)
	}

	if err := vm.Connect(); err != nil {
		proxy.Lock()
		delete(proxy.vms, hello.ContainerID)
		proxy.Unlock()
		response.SetError(err)
		return
	}

	client.vm = vm

	// We start one goroutine per-VM to monitor the qemu process
	proxy.wg.Add(1)
	go func() {
		<-vm.OnVMLost()
		vm.Close()
		proxy.wg.Done()
	}()
}

// "attach"
func attachHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	proxy := client.proxy

	attach := api.Attach{}
	if err := json.Unmarshal(data, &attach); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[attach.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", attach.ContainerID)
		return
	}

	client.infof(1, "attach(containerId=%s)", attach.ContainerID)

	client.vm = vm
}

// "bye"
func byeHandler(data []byte, userData interface{}, response *handlerResponse) {
	// Bye only affects the proxy.vms map and so removes the VM from the
	// client visible API.
	// vm.Close(), which tears down the VM object, is done at the end of
	// the VM life cycle, when  we detect the qemu process is effectively
	// gone (see helloHandler)

	client := userData.(*client)
	proxy := client.proxy

	bye := api.Bye{}
	if err := json.Unmarshal(data, &bye); err != nil {
		response.SetError(err)
		return
	}

	proxy.Lock()
	vm := proxy.vms[bye.ContainerID]
	proxy.Unlock()

	if vm == nil {
		response.SetErrorf("unknown containerID: %s", bye.ContainerID)
		return
	}

	client.info(1, "bye()")

	proxy.Lock()
	delete(proxy.vms, vm.containerID)
	proxy.Unlock()

	client.vm = nil
}

// "allocateIO"
func allocateIoHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	vm := client.vm

	allocateIo := api.AllocateIo{}
	if err := json.Unmarshal(data, &allocateIo); err != nil {
		response.SetError(err)
		return
	}

	if allocateIo.NStreams < 1 || allocateIo.NStreams > 2 {
		response.SetErrorf("asking for unexpected number of streams (%d)",
			allocateIo.NStreams)
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "allocateIo(nStreams=%d)", allocateIo.NStreams)

	// We'll send c0 to the client, keep c1
	c0, c1, err := Socketpair()
	if err != nil {
		response.SetError(err)
		return
	}

	f0, err := c0.File()
	if err != nil {
		response.SetError(err)
		return
	}

	ioBase := vm.AllocateIo(allocateIo.NStreams, client.id, c1)

	client.infof(1, "-> %d streams allocated, ioBase=%d", allocateIo.NStreams, ioBase)

	response.AddResult("ioBase", ioBase)
	response.SetFile(f0)

	// File() dups the underlying fd, so it's safe to close c0 here (will
	// keep the c0 <-> c1 connection alive).
	c0.Close()
}

// "hyper"
func hyperHandler(data []byte, userData interface{}, response *handlerResponse) {
	client := userData.(*client)
	hyper := api.Hyper{}
	vm := client.vm

	if err := json.Unmarshal(data, &hyper); err != nil {
		response.SetError(err)
		return
	}

	if vm == nil {
		response.SetErrorMsg("client not attached to a vm")
		return
	}

	client.infof(1, "hyper(cmd=%s, data=%s)", hyper.HyperName, hyper.Data)

	err := vm.SendMessage(hyper.HyperName, hyper.Data)
	response.SetError(err)
}

func newProxy() *proxy {
	return &proxy{
		vms: make(map[string]*vm),
	}
}

// DefaultSocketPath is populated at link time with the value of:
//   ${locatestatedir}/run/cc-oci-runtime/proxy
var DefaultSocketPath string

// ArgSocketPath is populated at runtime from the option -socket-path
var ArgSocketPath = flag.String("socket-path", "", "specify path to socket file")

func (proxy *proxy) init() error {
	var l net.Listener
	var err error

	// flags
	v := flag.Lookup("v").Value.(flag.Getter).Get().(glog.Level)
	proxy.enableVMConsole = v >= 3

	// Open the proxy socket
	fds := listenFds()

	if len(fds) > 1 {
		return fmt.Errorf("too many activated sockets (%d)", len(fds))
	} else if len(fds) == 1 {
		fd := fds[0]
		l, err = net.FileListener(fd)
		if err != nil {
			return fmt.Errorf("couldn't listen on socket: %v", err)
		}
	} else {
		// Invoking "go build" without any linker option will not
		// populate DefaultSocketPath, so fallback to a reasonable
		// path.
		if DefaultSocketPath == "" {
			DefaultSocketPath = "/var/run/cc-oci-runtime/proxy.sock"
		}

		socketPath := DefaultSocketPath
		if len(*ArgSocketPath) != 0 {
			socketPath = *ArgSocketPath
		}

		socketDir := filepath.Dir(socketPath)
		if err = os.MkdirAll(socketDir, 0750); err != nil {
			return fmt.Errorf("couldn't create socket directory: %v", err)
		}
		if err = os.Remove(socketPath); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("couldn't remove exiting socket: %v", err)
		}
		l, err = net.ListenUnix("unix", &net.UnixAddr{Name: socketPath, Net: "unix"})
		if err != nil {
			return fmt.Errorf("couldn't create AF_UNIX socket: %v", err)
		}
		if err = os.Chmod(socketPath, 0660|os.ModeSocket); err != nil {
			return fmt.Errorf("couldn't set mode on socket: %v", err)
		}

		glog.V(1).Info("listening on ", socketPath)
	}

	proxy.listener = l

	return nil
}

var nextClientID = uint64(1)

func (proxy *proxy) serveNewClient(proto *protocol, newConn net.Conn) {
	newClient := &client{
		id:    nextClientID,
		proxy: proxy,
		conn:  newConn,
	}

	atomic.AddUint64(&nextClientID, 1)

	// Unfortunately it's hard to find out information on the peer
	// at the other end of a unix socket. We use a per-client ID to
	// identify connections.
	newClient.info(1, "client connected")

	if err := proto.Serve(newConn, newClient); err != nil && err != io.EOF {
		newClient.infof(1, "error serving client: %v", err)
	}

	newConn.Close()
	newClient.info(1, "connection closed")
}

func (proxy *proxy) serve() {

	// Define the client (runtime/shim) <-> proxy protocol
	proto := newProtocol()
	proto.Handle("hello", helloHandler)
	proto.Handle("attach", attachHandler)
	proto.Handle("bye", byeHandler)
	proto.Handle("allocateIO", allocateIoHandler)
	proto.Handle("hyper", hyperHandler)

	glog.V(1).Info("proxy started")

	for {
		conn, err := proxy.listener.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "couldn't accept connection:", err)
			continue
		}

		go proxy.serveNewClient(proto, conn)
	}
}

func proxyMain() {
	proxy := newProxy()
	if err := proxy.init(); err != nil {
		fmt.Fprintln(os.Stderr, "init:", err.Error())
		os.Exit(1)
	}
	proxy.serve()

	// Wait for all the goroutines started by helloHandler to finish.
	//
	// Not stricly necessary as:
	//   • currently proxy.serve() cannot return,
	//   • even if it was, the process is about to exit anyway...
	//
	// That said, this wait group is used in the tests to ensure proper
	// serialisation between runs of proxyMain()(see proxy/proxy_test.go).
	proxy.wg.Wait()
}

func initLogging() {
	// We print logs on stderr by default.
	flag.Set("logtostderr", "true")

	// It can be practical to use an environment variable to trigger a verbose output
	level := os.Getenv("CC_PROXY_LOG_LEVEL")
	if level != "" {
		flag.Set("v", level)
	}
}

type profiler struct {
	enabled bool
	host    string
	port    uint
}

func (p *profiler) setup() {
	if !p.enabled {
		return
	}

	addr := fmt.Sprintf("%s:%d", p.host, p.port)
	url := "http://" + addr + "/debug/pprof"
	glog.V(1).Info("pprof enabled on " + url)

	go func() {
		http.ListenAndServe(addr, nil)
	}()
}

func main() {
	var pprof profiler

	initLogging()

	flag.BoolVar(&pprof.enabled, "pprof", false,
		"enable pprof ")
	flag.StringVar(&pprof.host, "pprof-host", "localhost",
		"host the pprof server will be bound to")
	flag.UintVar(&pprof.port, "pprof-port", 6060,
		"port the pprof server will be bound to")

	flag.Parse()
	defer glog.Flush()

	pprof.setup()
	proxyMain()
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	apps "k8s.io/api/apps/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/diff"
	kubeinformers "k8s.io/client-go/informers"
	k8sfake "k8s.io/client-go/kubernetes/fake"
	core "k8s.io/client-go/testing"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/record"

	samplecontroller "k8s.io/sample-controller/pkg/apis/samplecontroller/v1alpha1"
	"k8s.io/sample-controller/pkg/generated/clientset/versioned/fake"
	informers "k8s.io/sample-controller/pkg/generated/informers/externalversions"
)

var (
	alwaysReady        = func() bool { return true }
	noResyncPeriodFunc = func() time.Duration { return 0 }
)

type fixture struct {
	t *testing.T

	client     *fake.Clientset
	kubeclient *k8sfake.Clientset
	// Objects to put in the store.
	fooLister        []*samplecontroller.Foo
	deploymentLister []*apps.Deployment
	// Actions expected to happen on the client.
	kubeactions []core.Action
	actions     []core.Action
	// Objects from here preloaded into NewSimpleFake.
	kubeobjects []runtime.Object
	objects     []runtime.Object
}

func newFixture(t *testing.T) *fixture {
	f := &fixture{}
	f.t = t
	f.objects = []runtime.Object{}
	f.kubeobjects = []runtime.Object{}
	return f
}

func newFoo(name string, replicas *int32) *samplecontroller.Foo {
	return &samplecontroller.Foo{
		TypeMeta: metav1.TypeMeta{APIVersion: samplecontroller.SchemeGroupVersion.String()},
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metav1.NamespaceDefault,
		},
		Spec: samplecontroller.FooSpec{
			DeploymentName: fmt.Sprintf("%s-deployment", name),
			Replicas:       replicas,
		},
	}
}

func (f *fixture) newController() (*Controller, informers.SharedInformerFactory, kubeinformers.SharedInformerFactory) {
	f.client = fake.NewSimpleClientset(f.objects...)
	f.kubeclient = k8sfake.NewSimpleClientset(f.kubeobjects...)

	i := informers.NewSharedInformerFactory(f.client, noResyncPeriodFunc())
	k8sI := kubeinformers.NewSharedInformerFactory(f.kubeclient, noResyncPeriodFunc())

	c := NewController(f.kubeclient, f.client,
		k8sI.Apps().V1().Deployments(), i.Samplecontroller().V1alpha1().Foos())

	c.foosSynced = alwaysReady
	c.deploymentsSynced = alwaysReady
	c.recorder = &record.FakeRecorder{}

	for _, f := range f.fooLister {
		i.Samplecontroller().V1alpha1().Foos().Informer().GetIndexer().Add(f)
	}

	for _, d := range f.deploymentLister {
		k8sI.Apps().V1().Deployments().Informer().GetIndexer().Add(d)
	}

	return c, i, k8sI
}

func (f *fixture) run(fooName string) {
	f.runController(fooName, true, false)
}

func (f *fixture) runExpectError(fooName string) {
	f.runController(fooName, true, true)
}

func (f *fixture) runController(fooName string, startInformers bool, expectError bool) {
	c, i, k8sI := f.newController()
	if startInformers {
		stopCh := make(chan struct{})
		defer close(stopCh)
		i.Start(stopCh)
		k8sI.Start(stopCh)
	}

	err := c.syncHandler(fooName)
	if !expectError && err != nil {
		f.t.Errorf("error syncing foo: %v", err)
	} else if expectError && err == nil {
		f.t.Error("expected error syncing foo, got nil")
	}

	actions := filterInformerActions(f.client.Actions())
	for i, action := range actions {
		if len(f.actions) < i+1 {
			f.t.Errorf("%d unexpected actions: %+v", len(actions)-len(f.actions), actions[i:])
			break
		}

		expectedAction := f.actions[i]
		checkAction(expectedAction, action, f.t)
	}

	if len(f.actions) > len(actions) {
		f.t.Errorf("%d additional expected actions:%+v", len(f.actions)-len(actions), f.actions[len(actions):])
	}

	k8sActions := filterInformerActions(f.kubeclient.Actions())
	for i, action := range k8sActions {
		if len(f.kubeactions) < i+1 {
			f.t.Errorf("%d unexpected actions: %+v", len(k8sActions)-len(f.kubeactions), k8sActions[i:])
			break
		}

		expectedAction := f.kubeactions[i]
		checkAction(expectedAction, action, f.t)
	}

	if len(f.kubeactions) > len(k8sActions) {
		f.t.Errorf("%d additional expected actions:%+v", len(f.kubeactions)-len(k8sActions), f.kubeactions[len(k8sActions):])
	}
}

// checkAction verifies that expected and actual actions are equal and both have
// same attached resources
func checkAction(expected, actual core.Action, t *testing.T) {
	if !(expected.Matches(actual.GetVerb(), actual.GetResource().Resource) && actual.GetSubresource() == expected.GetSubresource()) {
		t.Errorf("Expected\n\t%#v\ngot\n\t%#v", expected, actual)
		return
	}

	if reflect.TypeOf(actual) != reflect.TypeOf(expected) {
		t.Errorf("Action has wrong type. Expected: %t. Got: %t", expected, actual)
		return
	}

	switch a := actual.(type) {
	case core.CreateAction:
		e, _ := expected.(core.CreateAction)
		expObject := e.GetObject()
		object := a.GetObject()

		if !reflect.DeepEqual(expObject, object) {
			t.Errorf("Action %s %s has wrong object\nDiff:\n %s",
				a.GetVerb(), a.GetResource().Resource, diff.ObjectGoPrintDiff(expObject, object))
		}
	case core.UpdateAction:
		e, _ := expected.(core.UpdateAction)
		expObject := e.GetObject()
		object := a.GetObject()

		if !reflect.DeepEqual(expObject, object) {
			t.Errorf("Action %s %s has wrong object\nDiff:\n %s",
				a.GetVerb(), a.GetResource().Resource, diff.ObjectGoPrintDiff(expObject, object))
		}
	case core.PatchAction:
		e, _ := expected.(core.PatchAction)
		expPatch := e.GetPatch()
		patch := a.GetPatch()

		if !reflect.DeepEqual(expPatch, patch) {
			t.Errorf("Action %s %s has wrong patch\nDiff:\n %s",
				a.GetVerb(), a.GetResource().Resource, diff.ObjectGoPrintDiff(expPatch, patch))
		}
	}
}

// filterInformerActions filters list and watch actions for testing resources.
// Since list and watch don't change resource state we can filter it to lower
// nose level in our tests.
func filterInformerActions(actions []core.Action) []core.Action {
	ret := []core.Action{}
	for _, action := range actions {
		if len(action.GetNamespace()) == 0 &&
			(action.Matches("list", "foos") ||
				action.Matches("watch", "foos") ||
				action.Matches("list", "deployments") ||
				action.Matches("watch", "deployments")) {
			continue
		}
		ret = append(ret, action)
	}

	return ret
}

func (f *fixture) expectCreateDeploymentAction(d *apps.Deployment) {
	f.kubeactions = append(f.kubeactions, core.NewCreateAction(schema.GroupVersionResource{Resource: "deployments"}, d.Namespace, d))
}

func (f *fixture) expectUpdateDeploymentAction(d *apps.Deployment) {
	f.kubeactions = append(f.kubeactions, core.NewUpdateAction(schema.GroupVersionResource{Resource: "deployments"}, d.Namespace, d))
}

func (f *fixture) expectUpdateFooStatusAction(foo *samplecontroller.Foo) {
	action := core.NewUpdateAction(schema.GroupVersionResource{Resource: "foos"}, foo.Namespace, foo)
	// TODO: Until #38113 is merged, we can't use Subresource
	//action.Subresource = "status"
	f.actions = append(f.actions, action)
}

func getKey(foo *samplecontroller.Foo, t *testing.T) string {
	key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(foo)
	if err != nil {
		t.Errorf("Unexpected error getting key for foo %v: %v", foo.Name, err)
		return ""
	}
	return key
}

func TestCreatesDeployment(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)

	expDeployment := newDeployment(foo)
	f.expectCreateDeploymentAction(expDeployment)
	f.expectUpdateFooStatusAction(foo)

	f.run(getKey(foo, t))
}

func TestDoNothing(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))
	d := newDeployment(foo)

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)
	f.deploymentLister = append(f.deploymentLister, d)
	f.kubeobjects = append(f.kubeobjects, d)

	f.expectUpdateFooStatusAction(foo)
	f.run(getKey(foo, t))
}

func TestUpdateDeployment(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))
	d := newDeployment(foo)

	// Update replicas
	foo.Spec.Replicas = int32Ptr(2)
	expDeployment := newDeployment(foo)

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)
	f.deploymentLister = append(f.deploymentLister, d)
	f.kubeobjects = append(f.kubeobjects, d)

	f.expectUpdateFooStatusAction(foo)
	f.expectUpdateDeploymentAction(expDeployment)
	f.run(getKey(foo, t))
}

func TestNotControlledByUs(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))
	d := newDeployment(foo)

	d.ObjectMeta.OwnerReferences = []metav1.OwnerReference{}

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)
	f.deploymentLister = append(f.deploymentLister, d)
	f.kubeobjects = append(f.kubeobjects, d)

	f.runExpectError(getKey(foo, t))
}

func int32Ptr(i int32) *int32 { return &i }
<commit_msg>Update deprecated diff.ObjectGoPrintDiff method<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	apps "k8s.io/api/apps/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/diff"
	kubeinformers "k8s.io/client-go/informers"
	k8sfake "k8s.io/client-go/kubernetes/fake"
	core "k8s.io/client-go/testing"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/record"

	samplecontroller "k8s.io/sample-controller/pkg/apis/samplecontroller/v1alpha1"
	"k8s.io/sample-controller/pkg/generated/clientset/versioned/fake"
	informers "k8s.io/sample-controller/pkg/generated/informers/externalversions"
)

var (
	alwaysReady        = func() bool { return true }
	noResyncPeriodFunc = func() time.Duration { return 0 }
)

type fixture struct {
	t *testing.T

	client     *fake.Clientset
	kubeclient *k8sfake.Clientset
	// Objects to put in the store.
	fooLister        []*samplecontroller.Foo
	deploymentLister []*apps.Deployment
	// Actions expected to happen on the client.
	kubeactions []core.Action
	actions     []core.Action
	// Objects from here preloaded into NewSimpleFake.
	kubeobjects []runtime.Object
	objects     []runtime.Object
}

func newFixture(t *testing.T) *fixture {
	f := &fixture{}
	f.t = t
	f.objects = []runtime.Object{}
	f.kubeobjects = []runtime.Object{}
	return f
}

func newFoo(name string, replicas *int32) *samplecontroller.Foo {
	return &samplecontroller.Foo{
		TypeMeta: metav1.TypeMeta{APIVersion: samplecontroller.SchemeGroupVersion.String()},
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metav1.NamespaceDefault,
		},
		Spec: samplecontroller.FooSpec{
			DeploymentName: fmt.Sprintf("%s-deployment", name),
			Replicas:       replicas,
		},
	}
}

func (f *fixture) newController() (*Controller, informers.SharedInformerFactory, kubeinformers.SharedInformerFactory) {
	f.client = fake.NewSimpleClientset(f.objects...)
	f.kubeclient = k8sfake.NewSimpleClientset(f.kubeobjects...)

	i := informers.NewSharedInformerFactory(f.client, noResyncPeriodFunc())
	k8sI := kubeinformers.NewSharedInformerFactory(f.kubeclient, noResyncPeriodFunc())

	c := NewController(f.kubeclient, f.client,
		k8sI.Apps().V1().Deployments(), i.Samplecontroller().V1alpha1().Foos())

	c.foosSynced = alwaysReady
	c.deploymentsSynced = alwaysReady
	c.recorder = &record.FakeRecorder{}

	for _, f := range f.fooLister {
		i.Samplecontroller().V1alpha1().Foos().Informer().GetIndexer().Add(f)
	}

	for _, d := range f.deploymentLister {
		k8sI.Apps().V1().Deployments().Informer().GetIndexer().Add(d)
	}

	return c, i, k8sI
}

func (f *fixture) run(fooName string) {
	f.runController(fooName, true, false)
}

func (f *fixture) runExpectError(fooName string) {
	f.runController(fooName, true, true)
}

func (f *fixture) runController(fooName string, startInformers bool, expectError bool) {
	c, i, k8sI := f.newController()
	if startInformers {
		stopCh := make(chan struct{})
		defer close(stopCh)
		i.Start(stopCh)
		k8sI.Start(stopCh)
	}

	err := c.syncHandler(fooName)
	if !expectError && err != nil {
		f.t.Errorf("error syncing foo: %v", err)
	} else if expectError && err == nil {
		f.t.Error("expected error syncing foo, got nil")
	}

	actions := filterInformerActions(f.client.Actions())
	for i, action := range actions {
		if len(f.actions) < i+1 {
			f.t.Errorf("%d unexpected actions: %+v", len(actions)-len(f.actions), actions[i:])
			break
		}

		expectedAction := f.actions[i]
		checkAction(expectedAction, action, f.t)
	}

	if len(f.actions) > len(actions) {
		f.t.Errorf("%d additional expected actions:%+v", len(f.actions)-len(actions), f.actions[len(actions):])
	}

	k8sActions := filterInformerActions(f.kubeclient.Actions())
	for i, action := range k8sActions {
		if len(f.kubeactions) < i+1 {
			f.t.Errorf("%d unexpected actions: %+v", len(k8sActions)-len(f.kubeactions), k8sActions[i:])
			break
		}

		expectedAction := f.kubeactions[i]
		checkAction(expectedAction, action, f.t)
	}

	if len(f.kubeactions) > len(k8sActions) {
		f.t.Errorf("%d additional expected actions:%+v", len(f.kubeactions)-len(k8sActions), f.kubeactions[len(k8sActions):])
	}
}

// checkAction verifies that expected and actual actions are equal and both have
// same attached resources
func checkAction(expected, actual core.Action, t *testing.T) {
	if !(expected.Matches(actual.GetVerb(), actual.GetResource().Resource) && actual.GetSubresource() == expected.GetSubresource()) {
		t.Errorf("Expected\n\t%#v\ngot\n\t%#v", expected, actual)
		return
	}

	if reflect.TypeOf(actual) != reflect.TypeOf(expected) {
		t.Errorf("Action has wrong type. Expected: %t. Got: %t", expected, actual)
		return
	}

	switch a := actual.(type) {
	case core.CreateAction:
		e, _ := expected.(core.CreateAction)
		expObject := e.GetObject()
		object := a.GetObject()

		if !reflect.DeepEqual(expObject, object) {
			t.Errorf("Action %s %s has wrong object\nDiff:\n %s",
				a.GetVerb(), a.GetResource().Resource, diff.ObjectGoPrintSideBySide(expObject, object))
		}
	case core.UpdateAction:
		e, _ := expected.(core.UpdateAction)
		expObject := e.GetObject()
		object := a.GetObject()

		if !reflect.DeepEqual(expObject, object) {
			t.Errorf("Action %s %s has wrong object\nDiff:\n %s",
				a.GetVerb(), a.GetResource().Resource, diff.ObjectGoPrintSideBySide(expObject, object))
		}
	case core.PatchAction:
		e, _ := expected.(core.PatchAction)
		expPatch := e.GetPatch()
		patch := a.GetPatch()

		if !reflect.DeepEqual(expPatch, patch) {
			t.Errorf("Action %s %s has wrong patch\nDiff:\n %s",
				a.GetVerb(), a.GetResource().Resource, diff.ObjectGoPrintSideBySide(expPatch, patch))
		}
	}
}

// filterInformerActions filters list and watch actions for testing resources.
// Since list and watch don't change resource state we can filter it to lower
// nose level in our tests.
func filterInformerActions(actions []core.Action) []core.Action {
	ret := []core.Action{}
	for _, action := range actions {
		if len(action.GetNamespace()) == 0 &&
			(action.Matches("list", "foos") ||
				action.Matches("watch", "foos") ||
				action.Matches("list", "deployments") ||
				action.Matches("watch", "deployments")) {
			continue
		}
		ret = append(ret, action)
	}

	return ret
}

func (f *fixture) expectCreateDeploymentAction(d *apps.Deployment) {
	f.kubeactions = append(f.kubeactions, core.NewCreateAction(schema.GroupVersionResource{Resource: "deployments"}, d.Namespace, d))
}

func (f *fixture) expectUpdateDeploymentAction(d *apps.Deployment) {
	f.kubeactions = append(f.kubeactions, core.NewUpdateAction(schema.GroupVersionResource{Resource: "deployments"}, d.Namespace, d))
}

func (f *fixture) expectUpdateFooStatusAction(foo *samplecontroller.Foo) {
	action := core.NewUpdateAction(schema.GroupVersionResource{Resource: "foos"}, foo.Namespace, foo)
	// TODO: Until #38113 is merged, we can't use Subresource
	//action.Subresource = "status"
	f.actions = append(f.actions, action)
}

func getKey(foo *samplecontroller.Foo, t *testing.T) string {
	key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(foo)
	if err != nil {
		t.Errorf("Unexpected error getting key for foo %v: %v", foo.Name, err)
		return ""
	}
	return key
}

func TestCreatesDeployment(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)

	expDeployment := newDeployment(foo)
	f.expectCreateDeploymentAction(expDeployment)
	f.expectUpdateFooStatusAction(foo)

	f.run(getKey(foo, t))
}

func TestDoNothing(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))
	d := newDeployment(foo)

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)
	f.deploymentLister = append(f.deploymentLister, d)
	f.kubeobjects = append(f.kubeobjects, d)

	f.expectUpdateFooStatusAction(foo)
	f.run(getKey(foo, t))
}

func TestUpdateDeployment(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))
	d := newDeployment(foo)

	// Update replicas
	foo.Spec.Replicas = int32Ptr(2)
	expDeployment := newDeployment(foo)

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)
	f.deploymentLister = append(f.deploymentLister, d)
	f.kubeobjects = append(f.kubeobjects, d)

	f.expectUpdateFooStatusAction(foo)
	f.expectUpdateDeploymentAction(expDeployment)
	f.run(getKey(foo, t))
}

func TestNotControlledByUs(t *testing.T) {
	f := newFixture(t)
	foo := newFoo("test", int32Ptr(1))
	d := newDeployment(foo)

	d.ObjectMeta.OwnerReferences = []metav1.OwnerReference{}

	f.fooLister = append(f.fooLister, foo)
	f.objects = append(f.objects, foo)
	f.deploymentLister = append(f.deploymentLister, d)
	f.kubeobjects = append(f.kubeobjects, d)

	f.runExpectError(getKey(foo, t))
}

func int32Ptr(i int32) *int32 { return &i }
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"log"
	"net/url"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/valyala/bytebufferpool"
	"github.com/valyala/fasthttp"
	"github.com/xtrafrancyz/vk-proxy/bytefmt"
	"github.com/xtrafrancyz/vk-proxy/replacer"
)

const (
	readBufferSize = 8192

	byteBufferPoolSetupRounds = 42000   // == bytebufferpool.calibrateCallsThreshold
	byteBufferPoolSetupSize   = 2097152 // 2**21
)

func init() {
	// Настройка размера буфера для ответа
	r := fasthttp.Response{}
	for i := 0; i <= byteBufferPoolSetupRounds; i++ {
		r.SetBodyString("")
		if b := r.Body(); cap(b) != byteBufferPoolSetupSize {
			r.SwapBody(make([]byte, byteBufferPoolSetupSize))
		} else {
			r.SwapBody(b[:cap(b)])
		}
		r.ResetBody()
	}

	// Настройка размера реплейсера
	for i := 0; i <= byteBufferPoolSetupRounds; i++ {
		b := replacer.AcquireBuffer()
		if cap(b.B) != byteBufferPoolSetupSize {
			b.B = make([]byte, byteBufferPoolSetupSize)
		} else {
			b.B = b.B[:cap(b.B)]
		}
		replacer.ReleaseBuffer(b)
	}
}

var (
	gzip            = []byte("gzip")
	vkProxyName     = []byte("vk-proxy")
	serverHeader    = []byte("Server")
	setCookie       = []byte("Set-Cookie")
	acceptEncoding  = []byte("Accept-Encoding")
	contentEncoding = []byte("Content-Encoding")
)

type ProxyConfig struct {
	ReduceMemoryUsage bool
	BaseDomain        string
	BaseStaticDomain  string
	LogVerbosity      int
	GzipUpstream      bool
	FilterFeed        bool
}

type Proxy struct {
	client   *fasthttp.Client
	server   *fasthttp.Server
	replacer *replacer.Replacer
	tracker  *tracker
	config   ProxyConfig
}

func NewProxy(config ProxyConfig) *Proxy {
	p := &Proxy{
		client: &fasthttp.Client{
			Name:           "vk-proxy",
			ReadBufferSize: readBufferSize,
			TLSConfig:      &tls.Config{InsecureSkipVerify: true},
		},
		replacer: &replacer.Replacer{
			ProxyBaseDomain:   config.BaseDomain,
			ProxyStaticDomain: config.BaseStaticDomain,
		},
		tracker: &tracker{
			uniqueUsers: make(map[string]bool),
		},
		config: config,
	}
	p.server = &fasthttp.Server{
		Handler:           p.handleProxy,
		ReduceMemoryUsage: config.ReduceMemoryUsage,
		ReadBufferSize:    readBufferSize,
		Name:              "vk-proxy",
	}
	if p.config.LogVerbosity > 0 {
		p.tracker.start()
	}
	return p
}

func (p *Proxy) ListenTCP(host string) error {
	log.Printf("Starting server on http://%s", host)
	return p.server.ListenAndServe(host)
}

func (p *Proxy) ListenUnix(path string) error {
	log.Printf("Starting server on http://unix:%s", path)
	return p.server.ListenAndServeUNIX(path, 0777)
}

func (p *Proxy) handleProxy(ctx *fasthttp.RequestCtx) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic when proxying the request: %s%s", r, debug.Stack())
			ctx.Error("500 Internal Server Error", 500)
		}
	}()
	start := time.Now()

	replaceContext := &replacer.ReplaceContext{
		Method:     ctx.Method(),
		FilterFeed: p.config.FilterFeed,
	}

	if !p.prepareProxyRequest(ctx, replaceContext) {
		ctx.Error("400 Bad Request", 400)
		return
	}

	if replaceContext.Host == "api.vk.com" &&
		(replaceContext.Path == "/away" || replaceContext.Path == "/away.php") {
		p.handleAway(ctx)
		return
	}

	err := p.client.DoTimeout(&ctx.Request, &ctx.Response, 30*time.Second)
	if err == nil {
		err = p.processProxyResponse(ctx, replaceContext)
	}

	elapsed := time.Since(start).Round(100 * time.Microsecond)

	if err != nil {
		log.Printf("%s %s error: %s", elapsed, ctx.Path(), err)
		if strings.Contains(err.Error(), "timed out") || strings.Contains(err.Error(), "timeout") {
			ctx.Error("408 Request Timeout", 408)
		} else {
			ctx.Error("500 Internal Server Error", 500)
		}
		return
	}

	if p.config.LogVerbosity > 0 {
		ip := ctx.Request.Header.Peek("CF-Connecting-IP") // Cloudflare
		if ip == nil {
			ip = ctx.Request.Header.Peek("X-Real-IP") // nginx
		}
		if ip == nil {
			ip = []byte(ctx.RemoteIP().String()) // real
		}
		p.tracker.trackRequest(string(ip), len(ctx.Response.Body()))
	}

	if p.config.LogVerbosity == 2 {
		log.Printf("%s %s %s%s %s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))))
	} else if p.config.LogVerbosity == 3 {
		log.Printf("%s %s %s%s %s\n%s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))), ctx.Response.Body())
	}
}

func (p *Proxy) handleAway(ctx *fasthttp.RequestCtx) {
	to := string(ctx.QueryArgs().Peek("to"))
	if to == "" {
		ctx.Error("Bad Request: 'to' argument is not set", 400)
		return
	}
	to, err := url.QueryUnescape(to)
	if err != nil {
		ctx.Error("Bad Request: could not unescape url", 400)
		return
	}
	ctx.Redirect(to, fasthttp.StatusMovedPermanently)
}

func (p *Proxy) prepareProxyRequest(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) bool {
	// Routing
	req := &ctx.Request
	uri := string(req.RequestURI())
	host := ""
	if strings.HasPrefix(uri, "/@") {
		slashIndex := strings.IndexByte(uri[2:], '/')
		if slashIndex == -1 {
			return false
		}
		endpoint := uri[2 : slashIndex+2]
		if endpoint != "vk.com" &&
			!strings.HasSuffix(endpoint, ".vk.com") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.net") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.com") {
			return false
		}
		host = endpoint
		uri = uri[2+slashIndex:]
		req.SetRequestURI(uri)
	} else if altHost := req.Header.Peek("Proxy-Host"); altHost != nil {
		host = string(altHost)
		switch host {
		case "static.vk.com":
		case "oauth.vk.com":
		default:
			return false
		}
		req.Header.Del("Proxy-Host")
	} else {
		host = "api.vk.com"
	}
	req.SetHost(host)

	// Replace some request data
	replaceContext.Host = host
	replaceContext.Path = string(ctx.Path())
	p.replacer.DoReplaceRequest(req, replaceContext)

	// After req.URI() call it is impossible to modify URI
	req.URI().SetScheme("https")
	if p.config.GzipUpstream {
		req.Header.SetBytesKV(acceptEncoding, gzip)
	} else {
		req.Header.DelBytes(acceptEncoding)
	}
	return true
}

func (p *Proxy) processProxyResponse(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) error {
	res := &ctx.Response
	res.Header.DelBytes(setCookie)
	res.Header.SetBytesKV(serverHeader, vkProxyName)

	var buf *bytebufferpool.ByteBuffer
	// Gunzip body if needed
	if bytes.Contains(res.Header.PeekBytes(contentEncoding), gzip) {
		res.Header.DelBytes(contentEncoding)
		buf = replacer.AcquireBuffer()
		_, err := fasthttp.WriteGunzip(buf, res.Body())
		if err != nil {
			replacer.ReleaseBuffer(buf)
			return err
		}
		replacer.ReleaseBuffer(&bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		})
	} else {
		buf = &bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		}
	}

	buf = p.replacer.DoReplaceResponse(res, buf, replaceContext)

	// avoid copying and save old buffer
	buf.B = res.SwapBody(buf.B)
	if cap(buf.B) > 10 {
		replacer.ReleaseBuffer(buf)
	}
	return nil
}

type tracker struct {
	lock        sync.Mutex
	requests    uint32
	bytes       uint64
	uniqueUsers map[string]bool
}

func (t *tracker) start() {
	go func() {
		for range time.Tick(60 * time.Second) {
			t.lock.Lock()
			log.Printf("Requests: %d, Traffic: %s, Online: %d", t.requests, bytefmt.ByteSize(t.bytes), len(t.uniqueUsers))
			t.requests = 0
			t.bytes = 0
			t.uniqueUsers = make(map[string]bool)
			t.lock.Unlock()
		}
	}()
}

func (t *tracker) trackRequest(ip string, size int) {
	t.lock.Lock()

	t.uniqueUsers[ip] = true
	t.requests++
	t.bytes += uint64(size)

	t.lock.Unlock()
}
<commit_msg>Улучшено логирование ошибок<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"log"
	"net/url"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/valyala/bytebufferpool"
	"github.com/valyala/fasthttp"
	"github.com/xtrafrancyz/vk-proxy/bytefmt"
	"github.com/xtrafrancyz/vk-proxy/replacer"
)

const (
	readBufferSize = 8192

	byteBufferPoolSetupRounds = 42000   // == bytebufferpool.calibrateCallsThreshold
	byteBufferPoolSetupSize   = 2097152 // 2**21
)

func init() {
	// Настройка размера буфера для ответа
	r := fasthttp.Response{}
	for i := 0; i <= byteBufferPoolSetupRounds; i++ {
		r.SetBodyString("")
		if b := r.Body(); cap(b) != byteBufferPoolSetupSize {
			r.SwapBody(make([]byte, byteBufferPoolSetupSize))
		} else {
			r.SwapBody(b[:cap(b)])
		}
		r.ResetBody()
	}

	// Настройка размера реплейсера
	for i := 0; i <= byteBufferPoolSetupRounds; i++ {
		b := replacer.AcquireBuffer()
		if cap(b.B) != byteBufferPoolSetupSize {
			b.B = make([]byte, byteBufferPoolSetupSize)
		} else {
			b.B = b.B[:cap(b.B)]
		}
		replacer.ReleaseBuffer(b)
	}
}

var (
	gzip            = []byte("gzip")
	vkProxyName     = []byte("vk-proxy")
	serverHeader    = []byte("Server")
	setCookie       = []byte("Set-Cookie")
	acceptEncoding  = []byte("Accept-Encoding")
	contentEncoding = []byte("Content-Encoding")
)

type ProxyConfig struct {
	ReduceMemoryUsage bool
	BaseDomain        string
	BaseStaticDomain  string
	LogVerbosity      int
	GzipUpstream      bool
	FilterFeed        bool
}

type Proxy struct {
	client   *fasthttp.Client
	server   *fasthttp.Server
	replacer *replacer.Replacer
	tracker  *tracker
	config   ProxyConfig
}

func NewProxy(config ProxyConfig) *Proxy {
	p := &Proxy{
		client: &fasthttp.Client{
			Name:           "vk-proxy",
			ReadBufferSize: readBufferSize,
			TLSConfig:      &tls.Config{InsecureSkipVerify: true},
		},
		replacer: &replacer.Replacer{
			ProxyBaseDomain:   config.BaseDomain,
			ProxyStaticDomain: config.BaseStaticDomain,
		},
		tracker: &tracker{
			uniqueUsers: make(map[string]bool),
		},
		config: config,
	}
	p.server = &fasthttp.Server{
		Handler:           p.handleProxy,
		ReduceMemoryUsage: config.ReduceMemoryUsage,
		ReadBufferSize:    readBufferSize,
		Name:              "vk-proxy",
	}
	if p.config.LogVerbosity > 0 {
		p.tracker.start()
	}
	return p
}

func (p *Proxy) ListenTCP(host string) error {
	log.Printf("Starting server on http://%s", host)
	return p.server.ListenAndServe(host)
}

func (p *Proxy) ListenUnix(path string) error {
	log.Printf("Starting server on http://unix:%s", path)
	return p.server.ListenAndServeUNIX(path, 0777)
}

func (p *Proxy) handleProxy(ctx *fasthttp.RequestCtx) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic when proxying the request: %s%s", r, debug.Stack())
			ctx.Error("500 Internal Server Error", 500)
		}
	}()
	start := time.Now()

	replaceContext := &replacer.ReplaceContext{
		Method:     ctx.Method(),
		FilterFeed: p.config.FilterFeed,
	}

	if !p.prepareProxyRequest(ctx, replaceContext) {
		ctx.Error("400 Bad Request", 400)
		return
	}

	if replaceContext.Host == "api.vk.com" &&
		(replaceContext.Path == "/away" || replaceContext.Path == "/away.php") {
		p.handleAway(ctx)
		return
	}

	err := p.client.DoTimeout(&ctx.Request, &ctx.Response, 30*time.Second)
	if err == nil {
		err = p.processProxyResponse(ctx, replaceContext)
	}

	elapsed := time.Since(start).Round(100 * time.Microsecond)

	if err != nil {
		log.Printf("%s %s %s%s error: %s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(), err)
		if strings.Contains(err.Error(), "timed out") || strings.Contains(err.Error(), "timeout") {
			ctx.Error("408 Request Timeout", 408)
		} else {
			ctx.Error("500 Internal Server Error", 500)
		}
		return
	}

	if p.config.LogVerbosity > 0 {
		ip := ctx.Request.Header.Peek("CF-Connecting-IP") // Cloudflare
		if ip == nil {
			ip = ctx.Request.Header.Peek("X-Real-IP") // nginx
		}
		if ip == nil {
			ip = []byte(ctx.RemoteIP().String()) // real
		}
		p.tracker.trackRequest(string(ip), len(ctx.Response.Body()))
	}

	if p.config.LogVerbosity == 2 {
		log.Printf("%s %s %s%s %s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))))
	} else if p.config.LogVerbosity == 3 {
		log.Printf("%s %s %s%s %s\n%s", elapsed, ctx.Request.Header.Method(), ctx.Host(), ctx.Path(),
			bytefmt.ByteSize(uint64(len(ctx.Response.Body()))), ctx.Response.Body())
	}
}

func (p *Proxy) handleAway(ctx *fasthttp.RequestCtx) {
	to := string(ctx.QueryArgs().Peek("to"))
	if to == "" {
		ctx.Error("Bad Request: 'to' argument is not set", 400)
		return
	}
	to, err := url.QueryUnescape(to)
	if err != nil {
		ctx.Error("Bad Request: could not unescape url", 400)
		return
	}
	ctx.Redirect(to, fasthttp.StatusMovedPermanently)
}

func (p *Proxy) prepareProxyRequest(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) bool {
	// Routing
	req := &ctx.Request
	uri := string(req.RequestURI())
	host := ""
	if strings.HasPrefix(uri, "/@") {
		slashIndex := strings.IndexByte(uri[2:], '/')
		if slashIndex == -1 {
			return false
		}
		endpoint := uri[2 : slashIndex+2]
		if endpoint != "vk.com" &&
			!strings.HasSuffix(endpoint, ".vk.com") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.net") &&
			!strings.HasSuffix(endpoint, ".vkuseraudio.com") {
			return false
		}
		host = endpoint
		uri = uri[2+slashIndex:]
		req.SetRequestURI(uri)
	} else if altHost := req.Header.Peek("Proxy-Host"); altHost != nil {
		host = string(altHost)
		switch host {
		case "static.vk.com":
		case "oauth.vk.com":
		default:
			return false
		}
		req.Header.Del("Proxy-Host")
	} else {
		host = "api.vk.com"
	}
	req.SetHost(host)

	// Replace some request data
	replaceContext.Host = host
	replaceContext.Path = string(ctx.Path())
	p.replacer.DoReplaceRequest(req, replaceContext)

	// After req.URI() call it is impossible to modify URI
	req.URI().SetScheme("https")
	if p.config.GzipUpstream {
		req.Header.SetBytesKV(acceptEncoding, gzip)
	} else {
		req.Header.DelBytes(acceptEncoding)
	}
	return true
}

func (p *Proxy) processProxyResponse(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) error {
	res := &ctx.Response
	res.Header.DelBytes(setCookie)
	res.Header.SetBytesKV(serverHeader, vkProxyName)

	var buf *bytebufferpool.ByteBuffer
	// Gunzip body if needed
	if bytes.Contains(res.Header.PeekBytes(contentEncoding), gzip) {
		res.Header.DelBytes(contentEncoding)
		buf = replacer.AcquireBuffer()
		_, err := fasthttp.WriteGunzip(buf, res.Body())
		if err != nil {
			replacer.ReleaseBuffer(buf)
			return err
		}
		replacer.ReleaseBuffer(&bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		})
	} else {
		buf = &bytebufferpool.ByteBuffer{
			B: res.SwapBody(nil),
		}
	}

	buf = p.replacer.DoReplaceResponse(res, buf, replaceContext)

	// avoid copying and save old buffer
	buf.B = res.SwapBody(buf.B)
	if cap(buf.B) > 10 {
		replacer.ReleaseBuffer(buf)
	}
	return nil
}

type tracker struct {
	lock        sync.Mutex
	requests    uint32
	bytes       uint64
	uniqueUsers map[string]bool
}

func (t *tracker) start() {
	go func() {
		for range time.Tick(60 * time.Second) {
			t.lock.Lock()
			log.Printf("Requests: %d, Traffic: %s, Online: %d", t.requests, bytefmt.ByteSize(t.bytes), len(t.uniqueUsers))
			t.requests = 0
			t.bytes = 0
			t.uniqueUsers = make(map[string]bool)
			t.lock.Unlock()
		}
	}()
}

func (t *tracker) trackRequest(ip string, size int) {
	t.lock.Lock()

	t.uniqueUsers[ip] = true
	t.requests++
	t.bytes += uint64(size)

	t.lock.Unlock()
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"text/template"
	"time"

	ctx "github.com/gorilla/context"
	"github.com/gorilla/mux"
	"github.com/jordan-wright/gophish/auth"
	"github.com/jordan-wright/gophish/models"
	"github.com/jordan-wright/gophish/worker"
)

var Worker *worker.Worker

func init() {
	Worker = worker.New()
	go Worker.Start()
}

// API (/api) provides access to api documentation
func API(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		templates := template.New("template")
		_, err := templates.ParseFiles("templates/api-docs.html")
		if err != nil {
			fmt.Println(err)
		}
		template.Must(templates, err).ExecuteTemplate(w, "base", nil)
	}
}

// API (/api/reset) resets a user's API key
func API_Reset(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "POST":
		u := ctx.Get(r, "user").(models.User)
		u.ApiKey = auth.GenerateSecureKey()
		err := models.PutUser(&u)
		if err != nil {
			http.Error(w, "Error setting API Key", http.StatusInternalServerError)
		} else {
			writeJSON(w, []byte(u.ApiKey))
		}
	}
}

// API_Campaigns returns a list of campaigns if requested via GET.
// If requested via POST, API_Campaigns creates a new campaign and returns a reference to it.
func API_Campaigns(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		cs, err := models.GetCampaigns(ctx.Get(r, "user_id").(int64))
		if err != nil {
			fmt.Println(err)
		}
		cj, err := json.MarshalIndent(cs, "", "  ")
		if checkError(err, w, "Error looking up campaigns", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, cj)
	//POST: Create a new campaign and return it as JSON
	case r.Method == "POST":
		c := models.Campaign{}
		// Put the request into a campaign
		err := json.NewDecoder(r.Body).Decode(&c)
		if checkError(err, w, "Invalid Request", http.StatusBadRequest) {
			return
		}
		if m, ok := models.ValidateCampaign(&c); !ok {
			http.Error(w, "Error: "+m, http.StatusBadRequest)
			return
		}
		// Fill in the details
		c.CreatedDate = time.Now()
		c.CompletedDate = time.Time{}
		c.Status = models.QUEUED
		c.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PostCampaign(&c, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Cannot insert campaign into database", http.StatusInternalServerError) {
			return
		}
		Worker.Queue <- &c
		cj, err := json.MarshalIndent(c, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, cj)
	}
}

// API_Campaigns_Id returns details about the requested campaign. If the campaign is not
// valid, API_Campaigns_Id returns null.
func API_Campaigns_Id(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, _ := strconv.ParseInt(vars["id"], 0, 64)
	switch {
	case r.Method == "GET":
		c := models.Campaign{}
		c, err := models.GetCampaign(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No campaign found", http.StatusNotFound) {
			return
		}
		cj, err := json.MarshalIndent(c, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, cj)
	case r.Method == "DELETE":
		_, err := models.GetCampaign(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No campaign found", http.StatusNotFound) {
			return
		}
		err = models.DeleteCampaign(id)
		if checkError(err, w, "Error deleting campaign", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, []byte("{\"success\" : \"true\"}"))
	}
}

// API_Groups returns details about the requested group. If the campaign is not
// valid, API_Groups returns null.
// Example:
/*
POST	/api/groups
		{ "name" : "Test Group",
		  "targets" : [
		  {
		  	"email" : "test@example.com"
		  },
		  { "email" : test2@example.com"
		  }]
		}

RESULT { "name" : "Test Group",
		  "targets" : [
		  {
		  	"email" : "test@example.com"
		  },
		  { "email" : test2@example.com"
		  }]
		}
*/
func API_Groups(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		gs, err := models.GetGroups(ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Groups not found", http.StatusNotFound) {
			return
		}
		gj, err := json.MarshalIndent(gs, "", "  ")
		if checkError(err, w, "Error marshaling group information", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, gj)
	//POST: Create a new group and return it as JSON
	case r.Method == "POST":
		g := models.Group{}
		// Put the request into a group
		err := json.NewDecoder(r.Body).Decode(&g)
		if checkError(err, w, "Invalid Request", http.StatusBadRequest) {
			return
		}
		// Check to make sure targets were specified
		if len(g.Targets) == 0 {
			http.Error(w, "Error: No targets specified", http.StatusBadRequest)
			return
		}
		g.ModifiedDate = time.Now()
		g.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PostGroup(&g)
		if checkError(err, w, "Error inserting group", http.StatusInternalServerError) {
			return
		}
		gj, err := json.MarshalIndent(g, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, gj)
	}
}

// API_Groups_Id returns details about the requested campaign. If the campaign is not
// valid, API_Campaigns_Id returns null.
func API_Groups_Id(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, _ := strconv.ParseInt(vars["id"], 0, 64)
	switch {
	case r.Method == "GET":
		g, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		gj, err := json.MarshalIndent(g, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, gj)
	case r.Method == "DELETE":
		g, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		err = models.DeleteGroup(&g)
		if checkError(err, w, "Error deleting group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, []byte("{\"success\" : \"true\"}"))
	case r.Method == "PUT":
		_, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		g := models.Group{}
		err = json.NewDecoder(r.Body).Decode(&g)
		if g.Id != id {
			http.Error(w, "Error: /:id and group_id mismatch", http.StatusBadRequest)
			return
		}
		// Check to make sure targets were specified
		if len(g.Targets) == 0 {
			http.Error(w, "Error: No targets specified", http.StatusBadRequest)
			return
		}
		g.ModifiedDate = time.Now()
		g.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PutGroup(&g)
		if checkError(err, w, "Error updating group", http.StatusInternalServerError) {
			return
		}
		gj, err := json.MarshalIndent(g, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, gj)
	}
}

func API_Templates(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		ts, err := models.GetTemplates(ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Templates not found", http.StatusNotFound) {
			return
		}
		tj, err := json.MarshalIndent(ts, "", "  ")
		if checkError(err, w, "Error marshaling template information", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, tj)
	//POST: Create a new template and return it as JSON
	case r.Method == "POST":
		t := models.Template{}
		// Put the request into a template
		err := json.NewDecoder(r.Body).Decode(&t)
		if checkError(err, w, "Invalid Request", http.StatusBadRequest) {
			return
		}
		t.ModifiedDate = time.Now()
		t.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PostTemplate(&t)
		if checkError(err, w, "Error inserting template", http.StatusInternalServerError) {
			return
		}
		tj, err := json.MarshalIndent(t, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, tj)
	}
}

func API_Templates_Id(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, _ := strconv.ParseInt(vars["id"], 0, 64)
	switch {
	case r.Method == "GET":
		t, err := models.GetTemplate(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No template found", http.StatusNotFound) {
			return
		}
		tj, err := json.MarshalIndent(t, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, tj)
	case r.Method == "DELETE":
		err := models.DeleteTemplate(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Error deleting group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, []byte("{\"success\" : \"true\"}"))
	case r.Method == "PUT":
		_, err := models.GetTemplate(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		t := models.Template{}
		err = json.NewDecoder(r.Body).Decode(&t)
		if t.Id != id {
			http.Error(w, "Error: /:id and template_id mismatch", http.StatusBadRequest)
			return
		}
		t.ModifiedDate = time.Now()
		err = models.PutTemplate(&t, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Error updating group", http.StatusInternalServerError) {
			return
		}
		tj, err := json.MarshalIndent(t, "", "  ")
		if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, tj)
	}
}

func writeJSON(w http.ResponseWriter, c []byte) {
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, "%s", c)
}
<commit_msg>Moved JSON encoding to helper function because DRY.<commit_after>package controllers

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"text/template"
	"time"

	ctx "github.com/gorilla/context"
	"github.com/gorilla/mux"
	"github.com/jordan-wright/gophish/auth"
	"github.com/jordan-wright/gophish/models"
	"github.com/jordan-wright/gophish/worker"
)

var Worker *worker.Worker

func init() {
	Worker = worker.New()
	go Worker.Start()
}

// API (/api) provides access to api documentation
func API(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		templates := template.New("template")
		_, err := templates.ParseFiles("templates/api-docs.html")
		if err != nil {
			fmt.Println(err)
		}
		template.Must(templates, err).ExecuteTemplate(w, "base", nil)
	}
}

// API (/api/reset) resets a user's API key
func API_Reset(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "POST":
		u := ctx.Get(r, "user").(models.User)
		u.ApiKey = auth.GenerateSecureKey()
		err := models.PutUser(&u)
		if err != nil {
			http.Error(w, "Error setting API Key", http.StatusInternalServerError)
		} else {
			writeJSON(w, []byte(u.ApiKey))
		}
	}
}

// API_Campaigns returns a list of campaigns if requested via GET.
// If requested via POST, API_Campaigns creates a new campaign and returns a reference to it.
func API_Campaigns(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		cs, err := models.GetCampaigns(ctx.Get(r, "user_id").(int64))
		if err != nil {
			fmt.Println(err)
		}
		writeJSON(w, cs)
	//POST: Create a new campaign and return it as JSON
	case r.Method == "POST":
		c := models.Campaign{}
		// Put the request into a campaign
		err := json.NewDecoder(r.Body).Decode(&c)
		if checkError(err, w, "Invalid Request", http.StatusBadRequest) {
			return
		}
		if m, ok := models.ValidateCampaign(&c); !ok {
			http.Error(w, "Error: "+m, http.StatusBadRequest)
			return
		}
		// Fill in the details
		c.CreatedDate = time.Now()
		c.CompletedDate = time.Time{}
		c.Status = models.QUEUED
		c.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PostCampaign(&c, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Cannot insert campaign into database", http.StatusInternalServerError) {
			return
		}
		Worker.Queue <- &c
		writeJSON(w, c)
	}
}

// API_Campaigns_Id returns details about the requested campaign. If the campaign is not
// valid, API_Campaigns_Id returns null.
func API_Campaigns_Id(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, _ := strconv.ParseInt(vars["id"], 0, 64)
	switch {
	case r.Method == "GET":
		c := models.Campaign{}
		c, err := models.GetCampaign(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No campaign found", http.StatusNotFound) {
			return
		}
		writeJSON(w, c)
	case r.Method == "DELETE":
		_, err := models.GetCampaign(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No campaign found", http.StatusNotFound) {
			return
		}
		err = models.DeleteCampaign(id)
		if checkError(err, w, "Error deleting campaign", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, []byte("{\"success\" : \"true\"}"))
	}
}

// API_Groups returns details about the requested group. If the campaign is not
// valid, API_Groups returns null.
// Example:
/*
POST	/api/groups
		{ "name" : "Test Group",
		  "targets" : [
		  {
		  	"email" : "test@example.com"
		  },
		  { "email" : test2@example.com"
		  }]
		}

RESULT { "name" : "Test Group",
		  "targets" : [
		  {
		  	"email" : "test@example.com"
		  },
		  { "email" : test2@example.com"
		  }]
		}
*/
func API_Groups(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		gs, err := models.GetGroups(ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Groups not found", http.StatusNotFound) {
			return
		}
		writeJSON(w, gs)
	//POST: Create a new group and return it as JSON
	case r.Method == "POST":
		g := models.Group{}
		// Put the request into a group
		err := json.NewDecoder(r.Body).Decode(&g)
		if checkError(err, w, "Invalid Request", http.StatusBadRequest) {
			return
		}
		// Check to make sure targets were specified
		if len(g.Targets) == 0 {
			http.Error(w, "Error: No targets specified", http.StatusBadRequest)
			return
		}
		g.ModifiedDate = time.Now()
		g.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PostGroup(&g)
		if checkError(err, w, "Error inserting group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, g)
	}
}

// API_Groups_Id returns details about the requested campaign. If the campaign is not
// valid, API_Campaigns_Id returns null.
func API_Groups_Id(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, _ := strconv.ParseInt(vars["id"], 0, 64)
	switch {
	case r.Method == "GET":
		g, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		writeJSON(w, g)
	case r.Method == "DELETE":
		g, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		err = models.DeleteGroup(&g)
		if checkError(err, w, "Error deleting group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, []byte("{\"success\" : \"true\"}"))
	case r.Method == "PUT":
		_, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		g := models.Group{}
		err = json.NewDecoder(r.Body).Decode(&g)
		if g.Id != id {
			http.Error(w, "Error: /:id and group_id mismatch", http.StatusBadRequest)
			return
		}
		// Check to make sure targets were specified
		if len(g.Targets) == 0 {
			http.Error(w, "Error: No targets specified", http.StatusBadRequest)
			return
		}
		g.ModifiedDate = time.Now()
		g.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PutGroup(&g)
		if checkError(err, w, "Error updating group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, g)
	}
}

func API_Templates(w http.ResponseWriter, r *http.Request) {
	switch {
	case r.Method == "GET":
		ts, err := models.GetTemplates(ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Templates not found", http.StatusNotFound) {
			return
		}
		writeJSON(w, ts)
	//POST: Create a new template and return it as JSON
	case r.Method == "POST":
		t := models.Template{}
		// Put the request into a template
		err := json.NewDecoder(r.Body).Decode(&t)
		if checkError(err, w, "Invalid Request", http.StatusBadRequest) {
			return
		}
		t.ModifiedDate = time.Now()
		t.UserId = ctx.Get(r, "user_id").(int64)
		err = models.PostTemplate(&t)
		if checkError(err, w, "Error inserting template", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, t)
	}
}

func API_Templates_Id(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id, _ := strconv.ParseInt(vars["id"], 0, 64)
	switch {
	case r.Method == "GET":
		t, err := models.GetTemplate(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No template found", http.StatusNotFound) {
			return
		}
		writeJSON(w, t)
	case r.Method == "DELETE":
		err := models.DeleteTemplate(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Error deleting group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, []byte("{\"success\" : \"true\"}"))
	case r.Method == "PUT":
		_, err := models.GetTemplate(id, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "No group found", http.StatusNotFound) {
			return
		}
		t := models.Template{}
		err = json.NewDecoder(r.Body).Decode(&t)
		if t.Id != id {
			http.Error(w, "Error: /:id and template_id mismatch", http.StatusBadRequest)
			return
		}
		t.ModifiedDate = time.Now()
		err = models.PutTemplate(&t, ctx.Get(r, "user_id").(int64))
		if checkError(err, w, "Error updating group", http.StatusInternalServerError) {
			return
		}
		writeJSON(w, t)
	}
}

func writeJSON(w http.ResponseWriter, c interface{}) {
	cj, err := json.MarshalIndent(c, "", "  ")
	if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
		return
	}
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, "%s", cj)
}
<|endoftext|>


<commit_before>package atlas

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
)

const atlasURL = "https://atlas.hashicorp.com"
const userAgent = "HashiCorp Atlas Go Client v1"

// If this is set to true, verbose debug data will be output
var Debug = false

// ErrAuth is the error returned if a 401 is returned by an API request.
var ErrAuth = fmt.Errorf("authentication failed")

// ErrNotFound is the error returned if a 404 is returned by an API request.
var ErrNotFound = fmt.Errorf("resource not found")

// RailsError represents an error that was returned from the Rails server.
type RailsError struct {
	Errors []string `json:"errors"`
}

// Error collects all of the errors in the RailsError and returns a comma-
// separated list of the errors that were returned from the server.
func (re *RailsError) Error() string {
	return strings.Join(re.Errors, ", ")
}

// Client represents a single connection to a Atlas API endpoint.
type Client struct {
	// URL is the full endpoint address to the Atlas server including the
	// protocol, port, and path.
	URL *url.URL

	// Token is the Atlas authentication token
	Token string

	// HTTPClient is the underlying http client with which to make requests.
	HTTPClient *http.Client
}

// DefaultClient returns a client that connects to the Atlas API.
func DefaultClient() *Client {
	client, err := NewClient(atlasURL)
	if err != nil {
		panic(err)
	}

	return client
}

// NewClient creates a new Atlas Client from the given URL (as a string). If
// the URL cannot be parsed, an error is returned. The HTTPClient is set to
// http.DefaultClient, but this can be changed programmatically by setting
// client.HTTPClient. The user can also programmatically set the URL as a
// *url.URL.
func NewClient(urlString string) (*Client, error) {
	if len(urlString) == 0 {
		return nil, fmt.Errorf("client: missing url")
	}

	parsedURL, err := url.Parse(urlString)
	if err != nil {
		return nil, err
	}

	client := &Client{
		URL:   parsedURL,
		Token: os.Getenv("ATLAS_TOKEN"),
	}

	if err := client.init(); err != nil {
		return nil, err
	}

	return client, nil
}

// init() sets defaults on the client.
func (c *Client) init() error {
	c.HTTPClient = http.DefaultClient
	return nil
}

// RequestOptions is the list of options to pass to the request.
type RequestOptions struct {
	// Params is a map of key-value pairs that will be added to the Request.
	Params map[string]string

	// Headers is a map of key-value pairs that will be added to the Request.
	Headers map[string]string

	// Body is an io.Reader object that will be streamed or uploaded with the
	// Request. BodyLength is the final size of the Body.
	Body       io.Reader
	BodyLength int64
}

// Request creates a new HTTP request using the given verb and sub path.
func (c *Client) Request(verb, spath string, ro *RequestOptions) (*http.Request, error) {
	// Ensure we have a RequestOptions struct (passing nil is an acceptable)
	if ro == nil {
		ro = new(RequestOptions)
	}

	// Create a new URL with the appended path
	u := *c.URL
	u.Path = path.Join(c.URL.Path, spath)

	// Add the token and other params
	if c.Token != "" {
		if ro.Params == nil {
			ro.Params = make(map[string]string)
		}

		ro.Params["access_token"] = c.Token
	}

	return c.rawRequest(verb, &u, ro)
}

func (c *Client) putFile(rawURL string, r io.Reader, size int64) error {
	url, err := url.Parse(rawURL)
	if err != nil {
		return err
	}

	request, err := c.rawRequest("PUT", url, &RequestOptions{
		Body:       r,
		BodyLength: size,
	})
	if err != nil {
		return err
	}

	if _, err := checkResp(c.HTTPClient.Do(request)); err != nil {
		return err
	}

	return nil
}

// rawRequest accepts a verb, URL, and RequestOptions struct and returns the
// constructed http.Request and any errors that occurred
func (c *Client) rawRequest(verb string, u *url.URL, ro *RequestOptions) (*http.Request, error) {
	if verb == "" {
		return nil, fmt.Errorf("client: missing verb")
	}

	if u == nil {
		return nil, fmt.Errorf("client: missing URL.url")
	}

	if ro == nil {
		return nil, fmt.Errorf("client: missing RequestOptions")
	}

	// Add the token and other params
	var params = make(url.Values)
	for k, v := range ro.Params {
		params.Add(k, v)
	}
	u.RawQuery = params.Encode()

	// Create the request object
	request, err := http.NewRequest(verb, u.String(), ro.Body)
	if err != nil {
		return nil, err
	}

	// Set the User-Agent
	request.Header.Set("User-Agent", userAgent)

	// Add any headers
	for k, v := range ro.Headers {
		request.Header.Add(k, v)
	}

	// Add content-length if we have it
	if ro.BodyLength > 0 {
		request.ContentLength = ro.BodyLength
	}

	return request, nil
}

// checkResp wraps http.Client.Do() and verifies that the request was
// successful. A non-200 request returns an error formatted to included any
// validation problems or otherwise.
func checkResp(resp *http.Response, err error) (*http.Response, error) {
	// If the err is already there, there was an error higher
	// up the chain, so just return that
	if err != nil {
		return resp, err
	}

	switch resp.StatusCode {
	case 200:
		return resp, nil
	case 201:
		return resp, nil
	case 202:
		return resp, nil
	case 204:
		return resp, nil
	case 400:
		return nil, parseErr(resp)
	case 401:
		return nil, ErrAuth
	case 404:
		return nil, ErrNotFound
	case 422:
		return nil, parseErr(resp)
	default:
		return nil, fmt.Errorf("client: %s", resp.Status)
	}
}

// parseErr is used to take an error JSON response and return a single string
// for use in error messages.
func parseErr(r *http.Response) error {
	re := &RailsError{}

	if err := decodeJSON(r, &re); err != nil {
		return fmt.Errorf("error decoding JSON body: %s", err)
	}

	return re
}

// decodeJSON is used to JSON decode a body into an interface.
func decodeJSON(resp *http.Response, out interface{}) error {
	defer resp.Body.Close()

	var r io.Reader = resp.Body
	if Debug {
		var buf bytes.Buffer
		r = io.TeeReader(resp.Body, &buf)
		defer func() {
			log.Printf("[DEBUG] client: decoding: %s", buf.String())
		}()
	}

	dec := json.NewDecoder(r)
	return dec.Decode(out)
}
<commit_msg>Log when the token is used from the environment<commit_after>package atlas

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
)

const atlasURL = "https://atlas.hashicorp.com"
const userAgent = "HashiCorp Atlas Go Client v1"

// If this is set to true, verbose debug data will be output
var Debug = false

// ErrAuth is the error returned if a 401 is returned by an API request.
var ErrAuth = fmt.Errorf("authentication failed")

// ErrNotFound is the error returned if a 404 is returned by an API request.
var ErrNotFound = fmt.Errorf("resource not found")

// RailsError represents an error that was returned from the Rails server.
type RailsError struct {
	Errors []string `json:"errors"`
}

// Error collects all of the errors in the RailsError and returns a comma-
// separated list of the errors that were returned from the server.
func (re *RailsError) Error() string {
	return strings.Join(re.Errors, ", ")
}

// Client represents a single connection to a Atlas API endpoint.
type Client struct {
	// URL is the full endpoint address to the Atlas server including the
	// protocol, port, and path.
	URL *url.URL

	// Token is the Atlas authentication token
	Token string

	// HTTPClient is the underlying http client with which to make requests.
	HTTPClient *http.Client
}

// DefaultClient returns a client that connects to the Atlas API.
func DefaultClient() *Client {
	client, err := NewClient(atlasURL)
	if err != nil {
		panic(err)
	}

	return client
}

// NewClient creates a new Atlas Client from the given URL (as a string). If
// the URL cannot be parsed, an error is returned. The HTTPClient is set to
// http.DefaultClient, but this can be changed programmatically by setting
// client.HTTPClient. The user can also programmatically set the URL as a
// *url.URL.
func NewClient(urlString string) (*Client, error) {
	if len(urlString) == 0 {
		return nil, fmt.Errorf("client: missing url")
	}

	parsedURL, err := url.Parse(urlString)
	if err != nil {
		return nil, err
	}

	token := os.Getenv("ATLAS_TOKEN")
	if token != "" {
		log.Printf("[DEBUG] using ATLAS_TOKEN (%s)", maskString(token))
	}

	client := &Client{
		URL:   parsedURL,
		Token: token,
	}

	if err := client.init(); err != nil {
		return nil, err
	}

	return client, nil
}

// init() sets defaults on the client.
func (c *Client) init() error {
	c.HTTPClient = http.DefaultClient
	return nil
}

// RequestOptions is the list of options to pass to the request.
type RequestOptions struct {
	// Params is a map of key-value pairs that will be added to the Request.
	Params map[string]string

	// Headers is a map of key-value pairs that will be added to the Request.
	Headers map[string]string

	// Body is an io.Reader object that will be streamed or uploaded with the
	// Request. BodyLength is the final size of the Body.
	Body       io.Reader
	BodyLength int64
}

// Request creates a new HTTP request using the given verb and sub path.
func (c *Client) Request(verb, spath string, ro *RequestOptions) (*http.Request, error) {
	// Ensure we have a RequestOptions struct (passing nil is an acceptable)
	if ro == nil {
		ro = new(RequestOptions)
	}

	// Create a new URL with the appended path
	u := *c.URL
	u.Path = path.Join(c.URL.Path, spath)

	// Add the token and other params
	if c.Token != "" {
		if ro.Params == nil {
			ro.Params = make(map[string]string)
		}

		ro.Params["access_token"] = c.Token
	}

	return c.rawRequest(verb, &u, ro)
}

func (c *Client) putFile(rawURL string, r io.Reader, size int64) error {
	url, err := url.Parse(rawURL)
	if err != nil {
		return err
	}

	request, err := c.rawRequest("PUT", url, &RequestOptions{
		Body:       r,
		BodyLength: size,
	})
	if err != nil {
		return err
	}

	if _, err := checkResp(c.HTTPClient.Do(request)); err != nil {
		return err
	}

	return nil
}

// rawRequest accepts a verb, URL, and RequestOptions struct and returns the
// constructed http.Request and any errors that occurred
func (c *Client) rawRequest(verb string, u *url.URL, ro *RequestOptions) (*http.Request, error) {
	if verb == "" {
		return nil, fmt.Errorf("client: missing verb")
	}

	if u == nil {
		return nil, fmt.Errorf("client: missing URL.url")
	}

	if ro == nil {
		return nil, fmt.Errorf("client: missing RequestOptions")
	}

	// Add the token and other params
	var params = make(url.Values)
	for k, v := range ro.Params {
		params.Add(k, v)
	}
	u.RawQuery = params.Encode()

	// Create the request object
	request, err := http.NewRequest(verb, u.String(), ro.Body)
	if err != nil {
		return nil, err
	}

	// Set the User-Agent
	request.Header.Set("User-Agent", userAgent)

	// Add any headers
	for k, v := range ro.Headers {
		request.Header.Add(k, v)
	}

	// Add content-length if we have it
	if ro.BodyLength > 0 {
		request.ContentLength = ro.BodyLength
	}

	return request, nil
}

// checkResp wraps http.Client.Do() and verifies that the request was
// successful. A non-200 request returns an error formatted to included any
// validation problems or otherwise.
func checkResp(resp *http.Response, err error) (*http.Response, error) {
	// If the err is already there, there was an error higher
	// up the chain, so just return that
	if err != nil {
		return resp, err
	}

	switch resp.StatusCode {
	case 200:
		return resp, nil
	case 201:
		return resp, nil
	case 202:
		return resp, nil
	case 204:
		return resp, nil
	case 400:
		return nil, parseErr(resp)
	case 401:
		return nil, ErrAuth
	case 404:
		return nil, ErrNotFound
	case 422:
		return nil, parseErr(resp)
	default:
		return nil, fmt.Errorf("client: %s", resp.Status)
	}
}

// parseErr is used to take an error JSON response and return a single string
// for use in error messages.
func parseErr(r *http.Response) error {
	re := &RailsError{}

	if err := decodeJSON(r, &re); err != nil {
		return fmt.Errorf("error decoding JSON body: %s", err)
	}

	return re
}

// decodeJSON is used to JSON decode a body into an interface.
func decodeJSON(resp *http.Response, out interface{}) error {
	defer resp.Body.Close()

	var r io.Reader = resp.Body
	if Debug {
		var buf bytes.Buffer
		r = io.TeeReader(resp.Body, &buf)
		defer func() {
			log.Printf("[DEBUG] client: decoding: %s", buf.String())
		}()
	}

	dec := json.NewDecoder(r)
	return dec.Decode(out)
}
<|endoftext|>