<commit_before>// Package reference provides a general type to represent any way of referencing images within the registry.
// Its main purpose is to abstract tags and digests (content-addressable hash).
//
// Grammar
//
// 	reference                       := repository [ ":" tag ] [ "@" digest ]
//	name                            := [hostname '/'] component ['/' component]*
//	hostname                        := hostcomponent ['.' hostcomponent]* [':' port-number]
//	hostcomponent                   := /([a-z0-9]|[a-z0-9][a-z0-9-]*[a-z0-9])/
//	port-number                     := /[0-9]+/
//	component                       := alpha-numeric [separator alpha-numeric]*
// 	alpha-numeric                   := /[a-z0-9]+/
//	separator                       := /[_.]|__|[-]*/
//
//	tag                             := /[\w][\w.-]{0,127}/
//
//	digest                          := digest-algorithm ":" digest-hex
//	digest-algorithm                := digest-algorithm-component [ digest-algorithm-separator digest-algorithm-component ]
//	digest-algorithm-separator      := /[+.-_]/
//	digest-algorithm-component      := /[A-Za-z][A-Za-z0-9]*/
//	digest-hex                      := /[0-9a-fA-F]{32,}/ ; At least 128 bit digest value
package reference

import (
	"errors"
	"fmt"

	"github.com/docker/distribution/digest"
)

const (
	// NameTotalLengthMax is the maximum total number of characters in a repository name.
	NameTotalLengthMax = 255
)

var (
	// ErrReferenceInvalidFormat represents an error while trying to parse a string as a reference.
	ErrReferenceInvalidFormat = errors.New("invalid reference format")

	// ErrTagInvalidFormat represents an error while trying to parse a string as a tag.
	ErrTagInvalidFormat = errors.New("invalid tag format")

	// ErrDigestInvalidFormat represents an error while trying to parse a string as a tag.
	ErrDigestInvalidFormat = errors.New("invalid digest format")

	// ErrNameEmpty is returned for empty, invalid repository names.
	ErrNameEmpty = errors.New("repository name must have at least one component")

	// ErrNameTooLong is returned when a repository name is longer than NameTotalLengthMax.
	ErrNameTooLong = fmt.Errorf("repository name must not be more than %v characters", NameTotalLengthMax)
)

// Reference is an opaque object reference identifier that may include
// modifiers such as a hostname, name, tag, and digest.
type Reference interface {
	// String returns the full reference
	String() string
}

// Field provides a wrapper type for resolving correct reference types when
// working with encoding.
type Field struct {
	reference Reference
}

// AsField wraps a reference in a Field for encoding.
func AsField(reference Reference) Field {
	return Field{reference}
}

// Reference unwraps the reference type from the field to
// return the Reference object. This object should be
// of the appropriate type to further check for different
// reference types.
func (f Field) Reference() Reference {
	return f.reference
}

// MarshalText serializes the field to byte text which
// is the string of the reference.
func (f Field) MarshalText() (p []byte, err error) {
	return []byte(f.reference.String()), nil
}

// UnmarshalText parses text bytes by invoking the
// reference parser to ensure the appropriately
// typed reference object is wrapped by field.
func (f *Field) UnmarshalText(p []byte) error {
	r, err := Parse(string(p))
	if err != nil {
		return err
	}

	f.reference = r
	return nil
}

// Named is an object with a full name
type Named interface {
	Reference
	Name() string
}

// Tagged is an object which has a tag
type Tagged interface {
	Reference
	Tag() string
}

// NamedTagged is an object including a name and tag.
type NamedTagged interface {
	Named
	Tag() string
}

// Digested is an object which has a digest
// in which it can be referenced by
type Digested interface {
	Reference
	Digest() digest.Digest
}

// Canonical reference is an object with a fully unique
// name including a name with hostname and digest
type Canonical interface {
	Named
	Digest() digest.Digest
}

// SplitHostname splits a named reference into a
// hostname and name string. If no valid hostname is
// found, the hostname is empty and the full value
// is returned as name
func SplitHostname(named Named) (string, string) {
	name := named.Name()
	match := anchoredNameRegexp.FindStringSubmatch(name)
	if match == nil || len(match) != 3 {
		return "", name
	}
	return match[1], match[2]
}

// Parse parses s and returns a syntactically valid Reference.
// If an error was encountered it is returned, along with a nil Reference.
// NOTE: Parse will not handle short digests.
func Parse(s string) (Reference, error) {
	matches := ReferenceRegexp.FindStringSubmatch(s)
	if matches == nil {
		if s == "" {
			return nil, ErrNameEmpty
		}
		// TODO(dmcgowan): Provide more specific and helpful error
		return nil, ErrReferenceInvalidFormat
	}

	if len(matches[1]) > NameTotalLengthMax {
		return nil, ErrNameTooLong
	}

	ref := reference{
		name: matches[1],
		tag:  matches[2],
	}
	if matches[3] != "" {
		var err error
		ref.digest, err = digest.ParseDigest(matches[3])
		if err != nil {
			return nil, err
		}
	}

	r := getBestReferenceType(ref)
	if r == nil {
		return nil, ErrNameEmpty
	}

	return r, nil
}

// ParseNamed parses s and returns a syntactically valid reference implementing
// the Named interface. The reference must have a name, otherwise an error is
// returned.
// If an error was encountered it is returned, along with a nil Reference.
// NOTE: ParseNamed will not handle short digests.
func ParseNamed(s string) (Named, error) {
	ref, err := Parse(s)
	if err != nil {
		return nil, err
	}
	named, isNamed := ref.(Named)
	if !isNamed {
		return nil, fmt.Errorf("reference %s has no name", ref.String())
	}
	return named, nil
}

// WithName returns a named object representing the given string. If the input
// is invalid ErrReferenceInvalidFormat will be returned.
func WithName(name string) (Named, error) {
	if len(name) > NameTotalLengthMax {
		return nil, ErrNameTooLong
	}
	if !anchoredNameRegexp.MatchString(name) {
		return nil, ErrReferenceInvalidFormat
	}
	return repository(name), nil
}

// WithTag combines the name from "name" and the tag from "tag" to form a
// reference incorporating both the name and the tag.
func WithTag(name Named, tag string) (NamedTagged, error) {
	if !anchoredTagRegexp.MatchString(tag) {
		return nil, ErrTagInvalidFormat
	}
	return taggedReference{
		name: name.Name(),
		tag:  tag,
	}, nil
}

// WithDigest combines the name from "name" and the digest from "digest" to form
// a reference incorporating both the name and the digest.
func WithDigest(name Named, digest digest.Digest) (Canonical, error) {
	if !anchoredDigestRegexp.MatchString(digest.String()) {
		return nil, ErrDigestInvalidFormat
	}
	return canonicalReference{
		name:   name.Name(),
		digest: digest,
	}, nil
}

func getBestReferenceType(ref reference) Reference {
	if ref.name == "" {
		// Allow digest only references
		if ref.digest != "" {
			return digestReference(ref.digest)
		}
		return nil
	}
	if ref.tag == "" {
		if ref.digest != "" {
			return canonicalReference{
				name:   ref.name,
				digest: ref.digest,
			}
		}
		return repository(ref.name)
	}
	if ref.digest == "" {
		return taggedReference{
			name: ref.name,
			tag:  ref.tag,
		}
	}

	return ref
}

type reference struct {
	name   string
	tag    string
	digest digest.Digest
}

func (r reference) String() string {
	return r.name + ":" + r.tag + "@" + r.digest.String()
}

func (r reference) Name() string {
	return r.name
}

func (r reference) Tag() string {
	return r.tag
}

func (r reference) Digest() digest.Digest {
	return r.digest
}

type repository string

func (r repository) String() string {
	return string(r)
}

func (r repository) Name() string {
	return string(r)
}

type digestReference digest.Digest

func (d digestReference) String() string {
	return d.String()
}

func (d digestReference) Digest() digest.Digest {
	return digest.Digest(d)
}

type taggedReference struct {
	name string
	tag  string
}

func (t taggedReference) String() string {
	return t.name + ":" + t.tag
}

func (t taggedReference) Name() string {
	return t.name
}

func (t taggedReference) Tag() string {
	return t.tag
}

type canonicalReference struct {
	name   string
	digest digest.Digest
}

func (c canonicalReference) String() string {
	return c.name + "@" + c.digest.String()
}

func (c canonicalReference) Name() string {
	return c.name
}

func (c canonicalReference) Digest() digest.Digest {
	return c.digest
}
<commit_msg>Reference name rather than repository in reference grammar<commit_after>// Package reference provides a general type to represent any way of referencing images within the registry.
// Its main purpose is to abstract tags and digests (content-addressable hash).
//
// Grammar
//
// 	reference                       := name [ ":" tag ] [ "@" digest ]
//	name                            := [hostname '/'] component ['/' component]*
//	hostname                        := hostcomponent ['.' hostcomponent]* [':' port-number]
//	hostcomponent                   := /([a-z0-9]|[a-z0-9][a-z0-9-]*[a-z0-9])/
//	port-number                     := /[0-9]+/
//	component                       := alpha-numeric [separator alpha-numeric]*
// 	alpha-numeric                   := /[a-z0-9]+/
//	separator                       := /[_.]|__|[-]*/
//
//	tag                             := /[\w][\w.-]{0,127}/
//
//	digest                          := digest-algorithm ":" digest-hex
//	digest-algorithm                := digest-algorithm-component [ digest-algorithm-separator digest-algorithm-component ]
//	digest-algorithm-separator      := /[+.-_]/
//	digest-algorithm-component      := /[A-Za-z][A-Za-z0-9]*/
//	digest-hex                      := /[0-9a-fA-F]{32,}/ ; At least 128 bit digest value
package reference

import (
	"errors"
	"fmt"

	"github.com/docker/distribution/digest"
)

const (
	// NameTotalLengthMax is the maximum total number of characters in a repository name.
	NameTotalLengthMax = 255
)

var (
	// ErrReferenceInvalidFormat represents an error while trying to parse a string as a reference.
	ErrReferenceInvalidFormat = errors.New("invalid reference format")

	// ErrTagInvalidFormat represents an error while trying to parse a string as a tag.
	ErrTagInvalidFormat = errors.New("invalid tag format")

	// ErrDigestInvalidFormat represents an error while trying to parse a string as a tag.
	ErrDigestInvalidFormat = errors.New("invalid digest format")

	// ErrNameEmpty is returned for empty, invalid repository names.
	ErrNameEmpty = errors.New("repository name must have at least one component")

	// ErrNameTooLong is returned when a repository name is longer than NameTotalLengthMax.
	ErrNameTooLong = fmt.Errorf("repository name must not be more than %v characters", NameTotalLengthMax)
)

// Reference is an opaque object reference identifier that may include
// modifiers such as a hostname, name, tag, and digest.
type Reference interface {
	// String returns the full reference
	String() string
}

// Field provides a wrapper type for resolving correct reference types when
// working with encoding.
type Field struct {
	reference Reference
}

// AsField wraps a reference in a Field for encoding.
func AsField(reference Reference) Field {
	return Field{reference}
}

// Reference unwraps the reference type from the field to
// return the Reference object. This object should be
// of the appropriate type to further check for different
// reference types.
func (f Field) Reference() Reference {
	return f.reference
}

// MarshalText serializes the field to byte text which
// is the string of the reference.
func (f Field) MarshalText() (p []byte, err error) {
	return []byte(f.reference.String()), nil
}

// UnmarshalText parses text bytes by invoking the
// reference parser to ensure the appropriately
// typed reference object is wrapped by field.
func (f *Field) UnmarshalText(p []byte) error {
	r, err := Parse(string(p))
	if err != nil {
		return err
	}

	f.reference = r
	return nil
}

// Named is an object with a full name
type Named interface {
	Reference
	Name() string
}

// Tagged is an object which has a tag
type Tagged interface {
	Reference
	Tag() string
}

// NamedTagged is an object including a name and tag.
type NamedTagged interface {
	Named
	Tag() string
}

// Digested is an object which has a digest
// in which it can be referenced by
type Digested interface {
	Reference
	Digest() digest.Digest
}

// Canonical reference is an object with a fully unique
// name including a name with hostname and digest
type Canonical interface {
	Named
	Digest() digest.Digest
}

// SplitHostname splits a named reference into a
// hostname and name string. If no valid hostname is
// found, the hostname is empty and the full value
// is returned as name
func SplitHostname(named Named) (string, string) {
	name := named.Name()
	match := anchoredNameRegexp.FindStringSubmatch(name)
	if match == nil || len(match) != 3 {
		return "", name
	}
	return match[1], match[2]
}

// Parse parses s and returns a syntactically valid Reference.
// If an error was encountered it is returned, along with a nil Reference.
// NOTE: Parse will not handle short digests.
func Parse(s string) (Reference, error) {
	matches := ReferenceRegexp.FindStringSubmatch(s)
	if matches == nil {
		if s == "" {
			return nil, ErrNameEmpty
		}
		// TODO(dmcgowan): Provide more specific and helpful error
		return nil, ErrReferenceInvalidFormat
	}

	if len(matches[1]) > NameTotalLengthMax {
		return nil, ErrNameTooLong
	}

	ref := reference{
		name: matches[1],
		tag:  matches[2],
	}
	if matches[3] != "" {
		var err error
		ref.digest, err = digest.ParseDigest(matches[3])
		if err != nil {
			return nil, err
		}
	}

	r := getBestReferenceType(ref)
	if r == nil {
		return nil, ErrNameEmpty
	}

	return r, nil
}

// ParseNamed parses s and returns a syntactically valid reference implementing
// the Named interface. The reference must have a name, otherwise an error is
// returned.
// If an error was encountered it is returned, along with a nil Reference.
// NOTE: ParseNamed will not handle short digests.
func ParseNamed(s string) (Named, error) {
	ref, err := Parse(s)
	if err != nil {
		return nil, err
	}
	named, isNamed := ref.(Named)
	if !isNamed {
		return nil, fmt.Errorf("reference %s has no name", ref.String())
	}
	return named, nil
}

// WithName returns a named object representing the given string. If the input
// is invalid ErrReferenceInvalidFormat will be returned.
func WithName(name string) (Named, error) {
	if len(name) > NameTotalLengthMax {
		return nil, ErrNameTooLong
	}
	if !anchoredNameRegexp.MatchString(name) {
		return nil, ErrReferenceInvalidFormat
	}
	return repository(name), nil
}

// WithTag combines the name from "name" and the tag from "tag" to form a
// reference incorporating both the name and the tag.
func WithTag(name Named, tag string) (NamedTagged, error) {
	if !anchoredTagRegexp.MatchString(tag) {
		return nil, ErrTagInvalidFormat
	}
	return taggedReference{
		name: name.Name(),
		tag:  tag,
	}, nil
}

// WithDigest combines the name from "name" and the digest from "digest" to form
// a reference incorporating both the name and the digest.
func WithDigest(name Named, digest digest.Digest) (Canonical, error) {
	if !anchoredDigestRegexp.MatchString(digest.String()) {
		return nil, ErrDigestInvalidFormat
	}
	return canonicalReference{
		name:   name.Name(),
		digest: digest,
	}, nil
}

func getBestReferenceType(ref reference) Reference {
	if ref.name == "" {
		// Allow digest only references
		if ref.digest != "" {
			return digestReference(ref.digest)
		}
		return nil
	}
	if ref.tag == "" {
		if ref.digest != "" {
			return canonicalReference{
				name:   ref.name,
				digest: ref.digest,
			}
		}
		return repository(ref.name)
	}
	if ref.digest == "" {
		return taggedReference{
			name: ref.name,
			tag:  ref.tag,
		}
	}

	return ref
}

type reference struct {
	name   string
	tag    string
	digest digest.Digest
}

func (r reference) String() string {
	return r.name + ":" + r.tag + "@" + r.digest.String()
}

func (r reference) Name() string {
	return r.name
}

func (r reference) Tag() string {
	return r.tag
}

func (r reference) Digest() digest.Digest {
	return r.digest
}

type repository string

func (r repository) String() string {
	return string(r)
}

func (r repository) Name() string {
	return string(r)
}

type digestReference digest.Digest

func (d digestReference) String() string {
	return d.String()
}

func (d digestReference) Digest() digest.Digest {
	return digest.Digest(d)
}

type taggedReference struct {
	name string
	tag  string
}

func (t taggedReference) String() string {
	return t.name + ":" + t.tag
}

func (t taggedReference) Name() string {
	return t.name
}

func (t taggedReference) Tag() string {
	return t.tag
}

type canonicalReference struct {
	name   string
	digest digest.Digest
}

func (c canonicalReference) String() string {
	return c.name + "@" + c.digest.String()
}

func (c canonicalReference) Name() string {
	return c.name
}

func (c canonicalReference) Digest() digest.Digest {
	return c.digest
}
<|endoftext|>


<commit_before>package possum

import (
	"fmt"
	"net/http"
	"net/http/pprof"
)

func (mux *ServeMux) InitPProf(prefix string) {
	if prefix == "" {
		prefix = "/debug/pprof"
	}
	mux.Handle(fmt.Sprintf("%s/", prefix), http.HandlerFunc(pprof.Index))
	mux.Handle(fmt.Sprintf("%s/cmdline", prefix), http.HandlerFunc(pprof.Cmdline))
	mux.Handle(fmt.Sprintf("%s/profile", prefix), http.HandlerFunc(pprof.Profile))
	mux.Handle(fmt.Sprintf("%s/symbol", prefix), http.HandlerFunc(pprof.Symbol))
}
<commit_msg>fixed pprof url<commit_after>package possum

import (
	"fmt"
	"log"
	"net/http"
	"net/http/pprof"
	rpprof "runtime/pprof"
	"strconv"
	"strings"
	"text/template"
)

func (mux *ServeMux) InitPProf(prefix string) {
	if prefix == "" {
		prefix = "/debug/pprof"
	}
	mux.Handle(fmt.Sprintf("%s/", prefix), pprofIndex(prefix))
	mux.Handle(fmt.Sprintf("%s/cmdline", prefix), http.HandlerFunc(pprof.Cmdline))
	mux.Handle(fmt.Sprintf("%s/profile", prefix), http.HandlerFunc(pprof.Profile))
	mux.Handle(fmt.Sprintf("%s/symbol", prefix), http.HandlerFunc(pprof.Symbol))
}

const pprofTemp = `

<title>%[1]s/</title>

%[1]s/

profiles:

{{range .}}

{{.Count}}{{.Name}}
{{end}}

full goroutine stack dump

`

func pprofIndex(prefix string) http.HandlerFunc {

	var indexTmpl = template.Must(template.New("index").Parse(fmt.Sprintf(pprofTemp, prefix)))
	if prefix[len(prefix)-1] != '/' {
		prefix = fmt.Sprintf("%s/", prefix)
	}
	f := func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, prefix) {
			name := strings.TrimPrefix(r.URL.Path, prefix)
			if name != "" {
				w.Header().Set("Content-Type", "text/plain; charset=utf-8")
				debug, _ := strconv.Atoi(r.FormValue("debug"))
				p := rpprof.Lookup(string(name))
				if p == nil {
					w.WriteHeader(404)
					fmt.Fprintf(w, "Unknown profile: %s\n", name)
					return
				}
				p.WriteTo(w, debug)
				return
			}
		}

		profiles := rpprof.Profiles()
		if err := indexTmpl.Execute(w, profiles); err != nil {
			log.Print(err)
		}
	}
	return f
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"
)

// TODO: Prevent any "localhost" or rfc1918 requests to our networks

type Coordinator struct {
	mu        sync.Mutex
	waiting   map[string]chan *http.Request
	responses map[string]chan *http.Response
}

func NewCoordinator() *Coordinator {
	return &Coordinator{
		waiting:   map[string]chan *http.Request{},
		responses: map[string]chan *http.Response{},
	}
}

func (c *Coordinator) getRequestChannel(fqdn string) chan *http.Request {
	c.mu.Lock()
	defer c.mu.Unlock()
	ch, ok := c.waiting[fqdn]
	if !ok {
		ch = make(chan *http.Request)
		c.waiting[fqdn] = ch
	}
	return ch
}

func (c *Coordinator) getResponseChannel(id string) chan *http.Response {
	c.mu.Lock()
	defer c.mu.Unlock()
	ch, ok := c.responses[id]
	if !ok {
		ch = make(chan *http.Response)
		c.responses[id] = ch
	}
	return ch
}

// Request a scrape.
func (c *Coordinator) DoScrape(ctx context.Context, r *http.Request) (*http.Response, error) {
	log.Printf("DoScrape %q", r.URL.String())
	r.Header.Add("Id", r.URL.String())
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("Matching client not found for %q: %s", r.URL.String(), ctx.Err())
	case c.getRequestChannel(r.URL.Hostname()) <- r:
	}

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case resp := <-c.getResponseChannel(r.URL.String()):
		return resp, nil
	}
}

// Client registering to accept a scrape request. Blocking.
func (c *Coordinator) WaitForScrapeInstruction(fqdn string) (*http.Request, error) {
	log.Printf("WaitForScrapeInstruction %q", fqdn)
	// TODO: What if the client times out?
	ch := c.getRequestChannel(fqdn)
	for {
		request := <-ch
		select {
		case <-request.Context().Done():
			// Request has timed out, get another one.
		default:
			return request, nil
		}
	}
}

// Client sending a scrape result in.
func (c *Coordinator) ScrapeResult(r *http.Response) {
	id := r.Header.Get("Id")
	log.Printf("ScrapeResult %q", id)
	r.Header.Del("Id")
	c.getResponseChannel(id) <- r
}

func copyHttpResponse(resp *http.Response, w http.ResponseWriter) {
	for k, v := range resp.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func main() {
	coordinator := NewCoordinator()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Proxy request
		if r.URL.Host != "" {
			ctx, _ := context.WithTimeout(r.Context(), time.Second*10)
			request := r.WithContext(ctx)
			request.RequestURI = ""

			resp, err := coordinator.DoScrape(ctx, request)
			if err != nil {
				log.Println(err)
				http.Error(w, fmt.Sprintf("Error scraping %q: %s", request.URL.String(), err.Error()), 500)
				return
			}
			defer resp.Body.Close()
			copyHttpResponse(resp, w)
			return
		}

		// Client registering and asking for scrapes.
		if r.URL.Path == "/poll" {
			fqdn, _ := ioutil.ReadAll(r.Body)
			request, _ := coordinator.WaitForScrapeInstruction(strings.TrimSpace(string(fqdn)))
			request.WriteProxy(w) // Send full request as the body of the response.
			log.Printf("Responded to /poll with %q", request.URL.String())
			return
		}

		// Scrape response from client.
		if r.URL.Path == "/push" {
			buf := &bytes.Buffer{}
			io.Copy(buf, r.Body)
			scrapeResult, _ := http.ReadResponse(bufio.NewReader(buf), nil)
			log.Printf("Got /push for %q", scrapeResult.Header.Get("Id"))
			coordinator.ScrapeResult(scrapeResult)
			return
		}

		http.Error(w, "404: Unknown path", 404)
	})

	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>Remove todo that's no longer relevant<commit_after>package main

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"
)

type Coordinator struct {
	mu        sync.Mutex
	waiting   map[string]chan *http.Request
	responses map[string]chan *http.Response
}

func NewCoordinator() *Coordinator {
	return &Coordinator{
		waiting:   map[string]chan *http.Request{},
		responses: map[string]chan *http.Response{},
	}
}

func (c *Coordinator) getRequestChannel(fqdn string) chan *http.Request {
	c.mu.Lock()
	defer c.mu.Unlock()
	ch, ok := c.waiting[fqdn]
	if !ok {
		ch = make(chan *http.Request)
		c.waiting[fqdn] = ch
	}
	return ch
}

func (c *Coordinator) getResponseChannel(id string) chan *http.Response {
	c.mu.Lock()
	defer c.mu.Unlock()
	ch, ok := c.responses[id]
	if !ok {
		ch = make(chan *http.Response)
		c.responses[id] = ch
	}
	return ch
}

// Request a scrape.
func (c *Coordinator) DoScrape(ctx context.Context, r *http.Request) (*http.Response, error) {
	log.Printf("DoScrape %q", r.URL.String())
	r.Header.Add("Id", r.URL.String())
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("Matching client not found for %q: %s", r.URL.String(), ctx.Err())
	case c.getRequestChannel(r.URL.Hostname()) <- r:
	}

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case resp := <-c.getResponseChannel(r.URL.String()):
		return resp, nil
	}
}

// Client registering to accept a scrape request. Blocking.
func (c *Coordinator) WaitForScrapeInstruction(fqdn string) (*http.Request, error) {
	log.Printf("WaitForScrapeInstruction %q", fqdn)
	// TODO: What if the client times out?
	ch := c.getRequestChannel(fqdn)
	for {
		request := <-ch
		select {
		case <-request.Context().Done():
			// Request has timed out, get another one.
		default:
			return request, nil
		}
	}
}

// Client sending a scrape result in.
func (c *Coordinator) ScrapeResult(r *http.Response) {
	id := r.Header.Get("Id")
	log.Printf("ScrapeResult %q", id)
	r.Header.Del("Id")
	c.getResponseChannel(id) <- r
}

func copyHttpResponse(resp *http.Response, w http.ResponseWriter) {
	for k, v := range resp.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func main() {
	coordinator := NewCoordinator()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Proxy request
		if r.URL.Host != "" {
			ctx, _ := context.WithTimeout(r.Context(), time.Second*10)
			request := r.WithContext(ctx)
			request.RequestURI = ""

			resp, err := coordinator.DoScrape(ctx, request)
			if err != nil {
				log.Println(err)
				http.Error(w, fmt.Sprintf("Error scraping %q: %s", request.URL.String(), err.Error()), 500)
				return
			}
			defer resp.Body.Close()
			copyHttpResponse(resp, w)
			return
		}

		// Client registering and asking for scrapes.
		if r.URL.Path == "/poll" {
			fqdn, _ := ioutil.ReadAll(r.Body)
			request, _ := coordinator.WaitForScrapeInstruction(strings.TrimSpace(string(fqdn)))
			request.WriteProxy(w) // Send full request as the body of the response.
			log.Printf("Responded to /poll with %q", request.URL.String())
			return
		}

		// Scrape response from client.
		if r.URL.Path == "/push" {
			buf := &bytes.Buffer{}
			io.Copy(buf, r.Body)
			scrapeResult, _ := http.ReadResponse(bufio.NewReader(buf), nil)
			log.Printf("Got /push for %q", scrapeResult.Header.Get("Id"))
			coordinator.ScrapeResult(scrapeResult)
			return
		}

		http.Error(w, "404: Unknown path", 404)
	})

	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"sync"
	"time"

	"github.com/abbot/go-http-auth"
)

const (
	browserStarted int = iota
	browserFailed
	seleniumError
)

const (
	pingPath  string = "/ping"
	errPath   string = "/err"
	routePath string = "/wd/hub/session"
	proxyPath string = routePath + "/"
	head      int    = len(proxyPath)
	tail      int    = head + 32
	sessPart  int    = 4 // /wd/hub/session/{various length session}
)

var (
	quota    map[string]Browsers = make(map[string]Browsers)
	routes   Routes              = make(Routes)
	num      uint64
	numLock  sync.Mutex
	confLock sync.RWMutex
)

type Routes map[string]*Host

type caps map[string]interface{}

func (c *caps) capability(k string) string {
	dc := (*c)["desiredCapabilities"]
	switch dc.(type) {
	case map[string]interface{}:
		v := dc.(map[string]interface{})
		switch v[k].(type) {
		case string:
			return v[k].(string)
		}
	}
	return ""
}

func (c *caps) browser() string {
	return c.capability("browserName")
}

func (c *caps) version() string {
	return c.capability("version")
}

func (h *Host) session(c caps) (map[string]interface{}, int) {
	b, _ := json.Marshal(c)
	resp, err := http.Post(h.sessionUrl(), "application/json", bytes.NewReader(b))
	if err != nil {
		return nil, seleniumError
	}
	var reply map[string]interface{}
	err = json.NewDecoder(resp.Body).Decode(&reply)
	if err != nil {
		return nil, seleniumError
	}
	if resp.StatusCode != http.StatusOK {
		return reply, browserFailed
	}
	return reply, browserStarted
}

func reply(w http.ResponseWriter, msg map[string]interface{}) {
	reply, _ := json.Marshal(msg)
	w.Header().Set("Content-Type", "application/json")
	w.Write(reply)
}

func serial() uint64 {
	numLock.Lock()
	defer numLock.Unlock()
	id := num
	num++
	return id
}

func info(r *http.Request) (user, remote string) {
	user = "unknown"
	if u, _, ok := r.BasicAuth(); ok {
		user = u
	}
	remote = r.Header.Get("X-Forwarded-For")
	if remote != "" {
		return
	}
	remote, _, _ = net.SplitHostPort(r.RemoteAddr)
	return
}

func fmtBrowser(browser, version string) string {
	if version != "" {
		return fmt.Sprintf("%s-%s", browser, version)
	}
	return browser
}

func browserErrMsg(js map[string]interface{}) string {
	if js == nil {
		return ""
	}
	val, ok := js["value"].(map[string]interface{})
	if !ok {
		return ""
	}
	msg, ok := val["message"].(string)
	if !ok {
		return ""
	}
	return msg
}

func jsonErrMsg(msg string) string {
	message := make(map[string]string)
	message["message"] = msg
	value := make(map[string]interface{})
	value["value"] = message
	value["status"] = 13
	result, _ := json.Marshal(value)
	return string(result)
}

func route(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	id := serial()
	user, remote := info(r)
	var c caps
	err := json.NewDecoder(r.Body).Decode(&c)
	if err != nil {
		http.Error(w, fmt.Sprintf("bad json format: %s", err.Error()), http.StatusBadRequest)
		log.Printf("[%d] [BAD_JSON] [%s] [%s] [%v]\n", id, user, remote, err)
		return
	}
	browser, version := c.browser(), c.version()
	if browser == "" {
		http.Error(w, "browser not set", http.StatusBadRequest)
		log.Printf("[%d] [BROWSER_NOT_SET] [%s] [%s]\n", id, user, remote)
		return
	}
	count := 0
loop:
	for {
		confLock.RLock()
		browsers := quota[user]
		hosts, version := browsers.find(browser, version)
		confLock.RUnlock()
		if len(hosts) == 0 {
			http.Error(w, fmt.Sprintf("unsupported browser: %s", fmtBrowser(browser, version)), http.StatusNotFound)
			log.Printf("[%d] [UNSUPPORTED_BROWSER] [%s] [%s] [%s]\n", id, user, remote, fmtBrowser(browser, version))
			return
		}
		for h, i := hosts.choose(); ; h, i = hosts.choose() {
			count++
			if h == nil {
				break loop
			}
			log.Printf("[%d] [SESSION_ATTEMPTED] [%s] [%s] [%s] [%s] [%d]\n", id, user, remote, fmtBrowser(browser, version), h.net(), count)
			excludes := make([]string, 0)
			resp, status := h.session(c)
			switch status {
			case browserStarted:
				sess := resp["sessionId"].(string)
				resp["sessionId"] = h.sum() + sess
				reply(w, resp)
				log.Printf("[%d] [%.2fs] [SESSION_CREATED] [%s] [%s] [%s] [%s] [%s] [%d]\n", id, float64(time.Now().Sub(start).Seconds()), user, remote, fmtBrowser(browser, version), h.net(), sess, count)
				return
			case browserFailed:
				hosts = append(hosts[:i], hosts[i+1:]...)
			case seleniumError:
				excludes = append(excludes, h.region)
				hosts, version = browsers.find(browser, version, excludes...)
			}
			log.Printf("[%d] [SESSION_FAILED] [%s] [%s] [%s] [%s] %s\n", id, user, remote, fmtBrowser(browser, version), h.net(), browserErrMsg(resp))
			if len(hosts) == 0 {
				break loop
			}
		}
	}
	http.Error(w, jsonErrMsg(fmt.Sprintf("cannot create session %s on any hosts after %d attempt(s)", fmtBrowser(browser, version), count)), http.StatusInternalServerError)
	log.Printf("[%d] [SESSION_NOT_CREATED] [%s] [%s] [%s]\n", id, user, remote, fmtBrowser(browser, version))
}

func proxy(r *http.Request) {
	_, remote := info(r)
	r.URL.Scheme = "http"
	if len(r.URL.Path) > tail {
		sum := r.URL.Path[head:tail]
		proxyPath := r.URL.Path[:head] + r.URL.Path[tail:]
		if h, ok := routes[sum]; ok {
			if body, err := ioutil.ReadAll(r.Body); err == nil {
				r.Body.Close()
				var msg map[string]interface{}
				if err := json.Unmarshal(body, &msg); err == nil {
					delete(msg, "sessionId")
					body, _ = json.Marshal(msg)
					r.ContentLength = int64(len(body))
				}
				r.Body = ioutil.NopCloser(bytes.NewReader(body))
			}
			r.URL.Host = h.net()
			r.URL.Path = proxyPath
			if r.Method == "DELETE" {
				sess := strings.Split(proxyPath, "/")[sessPart]
				log.Printf("[SESSION_DELETED] [%s] [%s] [%s]\n", remote, h.net(), sess)
			}
			return
		}
	}
	r.URL.Host = listen
	r.URL.Path = errPath
}

func ping(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Ok\n"))
}

func err(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "route not found", http.StatusNotFound)
}

func postOnly(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "POST" {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		handler(w, r)
	}
}

func readConfig(fn string, browsers *Browsers) error {
	file, err := ioutil.ReadFile(fn)
	if err != nil {
		return errors.New(fmt.Sprintf("error reading configuration file %s: %v", fn, err))
	}
	if err := xml.Unmarshal(file, browsers); err != nil {
		return errors.New(fmt.Sprintf("error parsing configuration file %s: %v", fn, err))
	}
	return nil
}

func appendRoutes(routes Routes, config *Browsers) Routes {
	for _, b := range config.Browsers {
		for _, v := range b.Versions {
			for _, r := range v.Regions {
				for i, h := range r.Hosts {
					r.Hosts[i].region = r.Name
					routes[h.sum()] = &r.Hosts[i]
				}
			}
		}
	}
	return routes
}

func requireBasicAuth(authenticator *auth.BasicAuth, handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
	return authenticator.Wrap(func(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
		handler(w, &r.Request)
	})
}

func mux() http.Handler {
	mux := http.NewServeMux()
	authenticator := auth.NewBasicAuthenticator(
		"Selenium Grid Router",
		auth.HtpasswdFileProvider(users),
	)
	mux.HandleFunc(pingPath, ping)
	mux.HandleFunc(errPath, err)
	mux.HandleFunc(routePath, requireBasicAuth(authenticator, postOnly(route)))
	mux.Handle(proxyPath, &httputil.ReverseProxy{Director: proxy})
	return mux
}
<commit_msg>Overriding browser version on prefix match<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"sync"
	"time"

	"github.com/abbot/go-http-auth"
)

const (
	browserStarted int = iota
	browserFailed
	seleniumError
)

const (
	pingPath  string = "/ping"
	errPath   string = "/err"
	routePath string = "/wd/hub/session"
	proxyPath string = routePath + "/"
	head      int    = len(proxyPath)
	tail      int    = head + 32
	sessPart  int    = 4 // /wd/hub/session/{various length session}
)

var (
	quota    map[string]Browsers = make(map[string]Browsers)
	routes   Routes              = make(Routes)
	num      uint64
	numLock  sync.Mutex
	confLock sync.RWMutex
)

type Routes map[string]*Host

type caps map[string]interface{}

func (c *caps) capability(k string) string {
	dc := (*c)["desiredCapabilities"]
	switch dc.(type) {
	case map[string]interface{}:
		v := dc.(map[string]interface{})
		switch v[k].(type) {
		case string:
			return v[k].(string)
		}
	}
	return ""
}

func (c *caps) setCapability(k string, v string) {
	(*c)["desiredCapabilities"].(map[string]interface{})[k] = v
}

func (c *caps) browser() string {
	return c.capability("browserName")
}

func (c *caps) version() string {
	return c.capability("version")
}

func (c *caps) setVersion(version string) {
	c.setCapability("version", version)
}

func (h *Host) session(c caps) (map[string]interface{}, int) {
	b, _ := json.Marshal(c)
	resp, err := http.Post(h.sessionUrl(), "application/json", bytes.NewReader(b))
	if err != nil {
		return nil, seleniumError
	}
	var reply map[string]interface{}
	err = json.NewDecoder(resp.Body).Decode(&reply)
	if err != nil {
		return nil, seleniumError
	}
	if resp.StatusCode != http.StatusOK {
		return reply, browserFailed
	}
	return reply, browserStarted
}

func reply(w http.ResponseWriter, msg map[string]interface{}) {
	reply, _ := json.Marshal(msg)
	w.Header().Set("Content-Type", "application/json")
	w.Write(reply)
}

func serial() uint64 {
	numLock.Lock()
	defer numLock.Unlock()
	id := num
	num++
	return id
}

func info(r *http.Request) (user, remote string) {
	user = "unknown"
	if u, _, ok := r.BasicAuth(); ok {
		user = u
	}
	remote = r.Header.Get("X-Forwarded-For")
	if remote != "" {
		return
	}
	remote, _, _ = net.SplitHostPort(r.RemoteAddr)
	return
}

func fmtBrowser(browser, version string) string {
	if version != "" {
		return fmt.Sprintf("%s-%s", browser, version)
	}
	return browser
}

func browserErrMsg(js map[string]interface{}) string {
	if js == nil {
		return ""
	}
	val, ok := js["value"].(map[string]interface{})
	if !ok {
		return ""
	}
	msg, ok := val["message"].(string)
	if !ok {
		return ""
	}
	return msg
}

func jsonErrMsg(msg string) string {
	message := make(map[string]string)
	message["message"] = msg
	value := make(map[string]interface{})
	value["value"] = message
	value["status"] = 13
	result, _ := json.Marshal(value)
	return string(result)
}

func route(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	id := serial()
	user, remote := info(r)
	var c caps
	err := json.NewDecoder(r.Body).Decode(&c)
	if err != nil {
		http.Error(w, fmt.Sprintf("bad json format: %s", err.Error()), http.StatusBadRequest)
		log.Printf("[%d] [BAD_JSON] [%s] [%s] [%v]\n", id, user, remote, err)
		return
	}
	browser, version := c.browser(), c.version()
	if browser == "" {
		http.Error(w, "browser not set", http.StatusBadRequest)
		log.Printf("[%d] [BROWSER_NOT_SET] [%s] [%s]\n", id, user, remote)
		return
	}
	count := 0
loop:
	for {
		confLock.RLock()
		browsers := quota[user]
		hosts, version := browsers.find(browser, version)
		confLock.RUnlock()
		if len(hosts) == 0 {
			http.Error(w, fmt.Sprintf("unsupported browser: %s", fmtBrowser(browser, version)), http.StatusNotFound)
			log.Printf("[%d] [UNSUPPORTED_BROWSER] [%s] [%s] [%s]\n", id, user, remote, fmtBrowser(browser, version))
			return
		}
		for h, i := hosts.choose(); ; h, i = hosts.choose() {
			count++
			if h == nil {
				break loop
			}
			log.Printf("[%d] [SESSION_ATTEMPTED] [%s] [%s] [%s] [%s] [%d]\n", id, user, remote, fmtBrowser(browser, version), h.net(), count)
			excludes := make([]string, 0)
			c.setVersion(version)
			resp, status := h.session(c)
			switch status {
			case browserStarted:
				sess := resp["sessionId"].(string)
				resp["sessionId"] = h.sum() + sess
				reply(w, resp)
				log.Printf("[%d] [%.2fs] [SESSION_CREATED] [%s] [%s] [%s] [%s] [%s] [%d]\n", id, float64(time.Now().Sub(start).Seconds()), user, remote, fmtBrowser(browser, version), h.net(), sess, count)
				return
			case browserFailed:
				hosts = append(hosts[:i], hosts[i+1:]...)
			case seleniumError:
				excludes = append(excludes, h.region)
				hosts, version = browsers.find(browser, version, excludes...)
			}
			log.Printf("[%d] [SESSION_FAILED] [%s] [%s] [%s] [%s] %s\n", id, user, remote, fmtBrowser(browser, version), h.net(), browserErrMsg(resp))
			if len(hosts) == 0 {
				break loop
			}
		}
	}
	http.Error(w, jsonErrMsg(fmt.Sprintf("cannot create session %s on any hosts after %d attempt(s)", fmtBrowser(browser, version), count)), http.StatusInternalServerError)
	log.Printf("[%d] [SESSION_NOT_CREATED] [%s] [%s] [%s]\n", id, user, remote, fmtBrowser(browser, version))
}

func proxy(r *http.Request) {
	_, remote := info(r)
	r.URL.Scheme = "http"
	if len(r.URL.Path) > tail {
		sum := r.URL.Path[head:tail]
		proxyPath := r.URL.Path[:head] + r.URL.Path[tail:]
		if h, ok := routes[sum]; ok {
			if body, err := ioutil.ReadAll(r.Body); err == nil {
				r.Body.Close()
				var msg map[string]interface{}
				if err := json.Unmarshal(body, &msg); err == nil {
					delete(msg, "sessionId")
					body, _ = json.Marshal(msg)
					r.ContentLength = int64(len(body))
				}
				r.Body = ioutil.NopCloser(bytes.NewReader(body))
			}
			r.URL.Host = h.net()
			r.URL.Path = proxyPath
			if r.Method == "DELETE" {
				sess := strings.Split(proxyPath, "/")[sessPart]
				log.Printf("[SESSION_DELETED] [%s] [%s] [%s]\n", remote, h.net(), sess)
			}
			return
		}
	}
	r.URL.Host = listen
	r.URL.Path = errPath
}

func ping(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Ok\n"))
}

func err(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "route not found", http.StatusNotFound)
}

func postOnly(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "POST" {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		handler(w, r)
	}
}

func readConfig(fn string, browsers *Browsers) error {
	file, err := ioutil.ReadFile(fn)
	if err != nil {
		return errors.New(fmt.Sprintf("error reading configuration file %s: %v", fn, err))
	}
	if err := xml.Unmarshal(file, browsers); err != nil {
		return errors.New(fmt.Sprintf("error parsing configuration file %s: %v", fn, err))
	}
	return nil
}

func appendRoutes(routes Routes, config *Browsers) Routes {
	for _, b := range config.Browsers {
		for _, v := range b.Versions {
			for _, r := range v.Regions {
				for i, h := range r.Hosts {
					r.Hosts[i].region = r.Name
					routes[h.sum()] = &r.Hosts[i]
				}
			}
		}
	}
	return routes
}

func requireBasicAuth(authenticator *auth.BasicAuth, handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
	return authenticator.Wrap(func(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
		handler(w, &r.Request)
	})
}

func mux() http.Handler {
	mux := http.NewServeMux()
	authenticator := auth.NewBasicAuthenticator(
		"Selenium Grid Router",
		auth.HtpasswdFileProvider(users),
	)
	mux.HandleFunc(pingPath, ping)
	mux.HandleFunc(errPath, err)
	mux.HandleFunc(routePath, requireBasicAuth(authenticator, postOnly(route)))
	mux.Handle(proxyPath, &httputil.ReverseProxy{Director: proxy})
	return mux
}
<|endoftext|>


<commit_before>package main

import (
	// Input/Output
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"

	// Webserver
	"net/http"

	// AWS
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

var (
	port                 string
	awsRegion, awsBucket string
	s3Session            *s3.S3
)

// Get an environment variable or use a default value if not set
func getEnvOrDefault(envName, defaultVal string, fatal bool) (envVal string) {
	envVal = os.Getenv(envName)
	if len(envVal) == 0 {
		if fatal {
			log.Fatal("Unable to start as env " + envName + " is not defined")
		}
		envVal = defaultVal
		log.Output(1, "Using default "+envName+": "+envVal)
	} else {
		log.Output(1, envName+": "+envVal)
	}
	return
}

// Get all the environment varables for this application
func getAllEnvVariables() {
	// Get the port that this webserver will be running upon
	port = getEnvOrDefault("PORT", "8000", false)

	// Get the AWS credentials
	awsRegion = getEnvOrDefault("AWS_REGION", "eu-west-1", false)
	awsBucket = getEnvOrDefault("AWS_BUCKET", "", true)
	getEnvOrDefault("AWS_ACCESS_KEY_ID", "", true)
	getEnvOrDefault("AWS_SECRET_ACCESS_KEY", "", true)
}

// Serve a request for a S3 file
func serveS3File(w http.ResponseWriter, r *http.Request) {
	var method = r.Method
	var path = r.URL.Path[1:] // Remove the / from the start of the URL

	// A file with no path cannot be served
	if path == "" {
		http.Error(w, "Path must be provided", 400)
		return
	}

	switch method {
	case "GET":
		serveGetS3File(path, w, r)
	case "PUT":
		servePutS3File(path, w, r)
	default:
		http.Error(w, "Method "+method+" not supported", 405)
	}
}

// Serve a GET request for a S3 file
func serveGetS3File(filePath string, w http.ResponseWriter, r *http.Request) {
	params := &s3.GetObjectInput{Bucket: aws.String(awsBucket), Key: aws.String(filePath)}
	resp, err := s3Session.GetObject(params)
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok {
			switch awserr.Code() {
			case "NoSuchKey":
				http.NotFound(w, r)
			default:
				http.Error(w, "An internal error occurred: "+awserr.Code()+" = "+awserr.Message(), 500)
			}
		} else {
			http.Error(w, "An internal error occurred: "+awserr.Message(), 500)
		}
		return
	}

	// File is ready to download
	io.Copy(w, resp.Body)
}

// Serve a PUT request for a S3 file
func servePutS3File(filePath string, w http.ResponseWriter, r *http.Request) {
	// Convert the uploaded body to a byte array TODO fix this for large sizes
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "An internal error occurred: "+fmt.Sprintf("%+v", err), 500)
	}

	params := &s3.PutObjectInput{Bucket: aws.String(awsBucket), Key: aws.String(filePath), Body: bytes.NewReader(b)}
	_, err = s3Session.PutObject(params)
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok {
			http.Error(w, "An internal error occurred: "+awserr.Code()+" = "+awserr.Message(), 500)
		} else {
			http.Error(w, "An internal error occurred: "+awserr.Message(), 500)
		}
		return
	}

	// File has been created TODO do not return a 201 if the file was updated
	http.Redirect(w, r, "/"+filePath, 201)
}

// Main method
func main() {
	// Set up all the environment variables
	getAllEnvVariables()

	s3Session = s3.New(session.New(), &aws.Config{Region: aws.String(awsRegion)})

	// Run the webserver
	http.HandleFunc("/", serveS3File)
	err := http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<commit_msg>Consolidate error handling<commit_after>package main

import (
	// Input/Output
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"os"

	// Webserver
	"net/http"

	// AWS
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

var (
	port                 string
	awsRegion, awsBucket string
	s3Session            *s3.S3
)

// Get an environment variable or use a default value if not set
func getEnvOrDefault(envName, defaultVal string, fatal bool) (envVal string) {
	envVal = os.Getenv(envName)
	if len(envVal) == 0 {
		if fatal {
			log.Fatal("Unable to start as env " + envName + " is not defined")
		}
		envVal = defaultVal
		log.Output(1, "Using default "+envName+": "+envVal)
	} else {
		log.Output(1, envName+": "+envVal)
	}
	return
}

// Get all the environment varables for this application
func getAllEnvVariables() {
	// Get the port that this webserver will be running upon
	port = getEnvOrDefault("PORT", "8000", false)

	// Get the AWS credentials
	awsRegion = getEnvOrDefault("AWS_REGION", "eu-west-1", false)
	awsBucket = getEnvOrDefault("AWS_BUCKET", "", true)
	getEnvOrDefault("AWS_ACCESS_KEY_ID", "", true)
	getEnvOrDefault("AWS_SECRET_ACCESS_KEY", "", true)
}

// Serve a request for a S3 file
func serveS3File(w http.ResponseWriter, r *http.Request) {
	var method = r.Method
	var path = r.URL.Path[1:] // Remove the / from the start of the URL

	// A file with no path cannot be served
	if path == "" {
		http.Error(w, "Path must be provided", 400)
		return
	}

	switch method {
	case "GET":
		serveGetS3File(path, w, r)
	case "PUT":
		servePutS3File(path, w, r)
	default:
		http.Error(w, "Method "+method+" not supported", 405)
	}
}

// Serve a GET request for a S3 file
func serveGetS3File(filePath string, w http.ResponseWriter, r *http.Request) {
	params := &s3.GetObjectInput{Bucket: aws.String(awsBucket), Key: aws.String(filePath)}
	resp, err := s3Session.GetObject(params)

	if handleHttpException(w, err) != nil {
		return
	}

	// File is ready to download
	io.Copy(w, resp.Body)
}

// Serve a PUT request for a S3 file
func servePutS3File(filePath string, w http.ResponseWriter, r *http.Request) {
	// Convert the uploaded body to a byte array TODO fix this for large sizes
	b, err := ioutil.ReadAll(r.Body)

	if handleHttpException(w, err) != nil {
		return
	}

	params := &s3.PutObjectInput{Bucket: aws.String(awsBucket), Key: aws.String(filePath), Body: bytes.NewReader(b)}
	_, err = s3Session.PutObject(params)

	if handleHttpException(w, err) != nil {
		return
	}

	// File has been created TODO do not return a 201 if the file was updated
	http.Redirect(w, r, "/"+filePath, 201)
}

// Handle an exception and write to response
func handleHttpException(w http.ResponseWriter, err error) (e error) {
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok {
			// aws error
			switch awserr.Code() {
			case "NoSuchKey":
				http.Error(w, "Not found: "+awserr.Message(), 404)
			default:
				http.Error(w, "An internal error occurred: "+awserr.Code()+" = "+awserr.Message(), 500)
			}
		} else {
			// golang error
			http.Error(w, "An internal error occurred: "+err.Error(), 500)
		}
	}
	return err
}

// Main method
func main() {
	// Set up all the environment variables
	getAllEnvVariables()

	s3Session = s3.New(session.New(), &aws.Config{Region: aws.String(awsRegion)})

	// Run the webserver
	http.HandleFunc("/", serveS3File)
	err := http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"github.com/Sirupsen/logrus"
	"io"
	"net"
	"strconv"
	"sync"
)

const (
	NETBUF_SIZE = 2048 // bytes
)

//go:generate stringer -type KeepAliveModeType
type KeepAliveModeType int

const (
	KEEPALIVE_BOTH KeepAliveModeType = iota
	KEEPALIVE_NOBACKEND
)

const (
	KEEPALIVE_BOTH_STRING      = "both"
	KEEPALIVE_NOBACKEND_STRING = "nobackend"
)

var (
	poolNetBuffers sync.Pool
	keepAliveMode  KeepAliveModeType
)

// var-constants
var (
	HEAD_CONNECTION            = []byte("CONNECTION")
	HEAD_CONNECTION_KEEP_ALIVE = []byte("KEEP-ALIVE")
	HEAD_CONTENT_LENGTH        = []byte("CONTENT-LENGTH")
)

// Get or create network buffer for proxy
func netbufGet() (buf []byte) {
	bufInterface := poolNetBuffers.Get()
	if bufInterface == nil {
		buf = make([]byte, NETBUF_SIZE)
	} else {
		buf = bufInterface.([]byte)
		// prevent data leak
		for i := range buf {
			buf[i] = 0
		}

	}
	return buf
}

func netbufPut(buf []byte) {
	poolNetBuffers.Put(buf)
}

func proxyHTTPHeaders(cid ConnectionID, targetConn net.Conn, sourceConn net.Conn) (
	res struct {
		KeepAlive        bool
		ContentLength    int64
		HasContextLength bool
		Err              error
	}) {
	buf := netbufGet()
	defer netbufPut(buf)
	var totalReadBytes int

	// Read lines
readHeaderLines:
	for {
		var i int
		var headerStart []byte
		for i = 0; i < len(buf); i++ {
			var readBytes int
			readBytes, res.Err = sourceConn.Read(buf[i : i+1])
			totalReadBytes += readBytes
			if res.Err != nil {
				logrus.Debugf("Error while read header from '%v' cid '%v': %v", sourceConn.RemoteAddr(), cid, res.Err)
				return
			}
			if readBytes != 1 {
				logrus.Infof("Can't read a byte from header from '%v' cid '%v'", sourceConn.RemoteAddr(), cid)
				return
			}
			if buf[i] == ':' || buf[i] == '\n' {
				headerStart = buf[:i+1]
				logrus.Debugf("Header Name '%v' -> '%v' cid '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, buf[:i])
				break
			}
		}
		if len(headerStart) == 0 {
			logrus.Infof("Header line longer then buffer (%v bytes). Force close connection. '%v' -> '%v' cid '%v'.", len(buf), sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)
			return
		}

		// Empty line - end http headers
		if bytes.Equal(headerStart, []byte("\n")) || bytes.Equal(headerStart, []byte("\r\n")) {
			break readHeaderLines
		}

		headerName := headerStart[:len(headerStart)-1] // Cut trailing colon from start
		headerNameUpperCase := bytes.ToUpper(headerName)

		skipHeader := false
		for _, ownHeader := range cutHeaders {
			if bytes.Equal(ownHeader, headerNameUpperCase) {
				skipHeader = true
				break
			}
			ownHeaderS := string(ownHeader)
			headerNameS := string(headerNameUpperCase)
			skipHeader = ownHeaderS == headerNameS
		}

		if skipHeader {
			logrus.Debugf("Skip header: '%v' -> '%v' cid '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerNameUpperCase)
			buf[0] = headerStart[len(headerStart)-1]

			for buf[0] != '\n' {
				_, res.Err = sourceConn.Read(buf[:1])
				if res.Err != nil {
					logrus.Infof("Error read header. Close connections. '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, res.Err)
					sourceConn.Close()
					targetConn.Close()
					return
				}
			}
			continue readHeaderLines
		}

		logrus.Debugf("Copy header: '%v' -> '%v' cid '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerName)

		// copy header without changes
		_, res.Err = targetConn.Write(headerStart)
		if res.Err != nil {
			logrus.Infof("Write header start, from '%v' to '%v' cid '%v', headerStart='%s': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerStart, res.Err)
			return
		}

		needHeaderContent := bytes.Equal(headerNameUpperCase, HEAD_CONTENT_LENGTH) || bytes.Equal(headerNameUpperCase, HEAD_CONNECTION)
		headerContent := bytes.NewBuffer(buf[1:])
		headerContent.Reset()

		for buf[0] != '\n' {
			var readBytes int
			readBytes, res.Err = sourceConn.Read(buf[:1])
			if res.Err != nil {
				logrus.Infof("Error read header to copy. Close connections. '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, res.Err)
				return
			}
			if readBytes != 1 {
				logrus.Infof("Header copy read bytes != 1. Error. Close connections. '%v' -> '%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)
				return
			}
			_, res.Err = targetConn.Write(buf[:1])
			if res.Err != nil {
				logrus.Infof("Error write header. Close connections. '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, res.Err)
				return
			}
			if needHeaderContent {
				headerContent.WriteByte(buf[0])
			}
		}
		if needHeaderContent {
			switch {
			case bytes.Equal(headerNameUpperCase, HEAD_CONNECTION):
				res.KeepAlive = bytes.EqualFold(HEAD_CONNECTION_KEEP_ALIVE, bytes.TrimSpace(headerContent.Bytes()))

			case bytes.Equal(headerNameUpperCase, HEAD_CONTENT_LENGTH):
				res.ContentLength, res.Err = strconv.ParseInt(string(bytes.TrimSpace(headerContent.Bytes())), 10, 64)
				if res.Err == nil {
					res.HasContextLength = true
					logrus.Debugf("Header content-length parsed from '%v' to '%v' cid '%v': %v", sourceConn.RemoteAddr(),
						targetConn.RemoteAddr(), cid, res.ContentLength)
				} else {
					logrus.Infof("Can't header content-length parsed from '%v' to '%v' cid '%v' content '%s': %v", sourceConn.RemoteAddr(),
						targetConn.RemoteAddr(), cid, headerContent.Bytes(), res.Err)
				}

			default:
				logrus.Debugf("ERROR. Unknow why i need header content. Code error. From '%v' to '%v' cid '%v', header name '%s'",
					sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerNameUpperCase,
				)
			}
		}
	}

	remoteTcpAddr := sourceConn.RemoteAddr().(*net.TCPAddr)
	remoteAddrString := remoteTcpAddr.IP.String()

	headerBuf := bytes.NewBuffer(buf)
	headerBuf.Reset()

	// Write real IP
	for _, header := range realIPHeaderNames {
		headerBuf.Write(header)
		headerBuf.WriteByte(':')
		headerBuf.WriteString(remoteAddrString)
		headerBuf.WriteString("\r\n")
	}

	// Write CID
	if *connectionIdHeader != "" {
		headerBuf.WriteString(*connectionIdHeader)
		headerBuf.WriteString(": ")
		headerBuf.WriteString(cid.String())
		headerBuf.WriteString("\r\n")
	}

	headerBuf.Write(additionalHeaders)
	headerBuf.Write([]byte("\r\n")) // end http headers
	logrus.Debugf("Add headers. '%v' -> '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), headerBuf.Bytes())

	_, res.Err = targetConn.Write(headerBuf.Bytes())
	if res.Err != nil {
		logrus.Infof("Error while write real ip headers to target '%v' -> '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), res.Err)
	}

	return
}

func startProxy(cid ConnectionID, targetAddr net.TCPAddr, in net.Conn) {
	targetConnCommon, err := net.DialTimeout("tcp", targetAddr.String(), *targetConnTimeout)
	if err != nil {
		logrus.Warnf("Can't connect to target '%v' cid '%v': %v", targetAddr.String(), cid, err)
		in.Close()
		return
	}

	targetConn, ok := targetConnCommon.(*net.TCPConn)
	targetConn.SetKeepAlive(true)
	targetConn.SetKeepAlivePeriod(*tcpKeepAliveInterval)

	if !ok {
		logrus.Errorf("Can't cast connection to tcp connection, target '%v' cid '%v'", targetAddr.String(), cid)
		return
	}

	switch *proxyMode {
	case "http":
		startProxyHTTP(cid, targetConn, in)
	case "tcp":
		startProxyTCP(cid, targetConn, in)
	default:
		logrus.Panicf("Unknow proxy mode cid '%v': %v", cid, *proxyMode)
	}
}

func startProxyHTTP(cid ConnectionID, targetConn net.Conn, sourceConn net.Conn) {
	defer targetConn.Close()
	defer sourceConn.Close()

	logrus.Debugf("Start http-proxy connection from '%v' to '%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)

	buf := netbufGet()
	defer netbufPut(buf)
	var receivedBytesCount int64
	for {
		state := proxyHTTPHeaders(cid, targetConn, sourceConn)
		keepAlive := state.KeepAlive
		if state.Err != nil {
			logrus.Debug("Cid '%v'. Can't read headers: %v", cid, state.Err)
			return
		}
		if state.HasContextLength {
			logrus.Debugf("Start keep-alieved proxy. '%v' -> '%v' cid '%v', content-length '%v'", sourceConn.RemoteAddr(),
				targetConn.RemoteAddr(), cid, state.ContentLength)

			//request
			bytesCopied, err := io.CopyBuffer(targetConn, io.LimitReader(sourceConn, state.ContentLength), buf)
			receivedBytesCount += bytesCopied

			if err == nil {
				logrus.Debugf("Connection chunk copied '%v' -> '%v' cid '%v', bytes transferred '%v' (%v), error: %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, bytesCopied, receivedBytesCount, err)
			} else {
				logrus.Debugf("Connection closed '%v' -> '%v' cid '%v', bytes transferred '%v' (%v), error: %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, bytesCopied, receivedBytesCount, err)
				return
			}

			// proxy answer
			answerState := proxyHTTPHeaders(cid, sourceConn, targetConn)
			keepAlive = keepAlive && answerState.KeepAlive
			var reader io.Reader

			if answerState.HasContextLength {
				reader = io.LimitReader(targetConn, answerState.ContentLength)
			} else {
				reader = targetConn
			}

			_, err = io.CopyBuffer(sourceConn, reader, buf)
			if err != nil {
				logrus.Debugf("Cid '%v'. Can't copy answer to customer: %v", cid, err)
				return
			}

		} else {
			// answer from server proxy without changes
			go func() {
				buf := netbufGet()
				defer netbufPut(buf)

				_, err := io.CopyBuffer(sourceConn, targetConn, buf)
				logrus.Debugf("Connection closed with error1 '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, err)
				sourceConn.Close()
				targetConn.Close()
			}()

			logrus.Debugf("Start proxy without support keepalive middle headers '%v' -> '%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)
			bytesCopied, err := io.CopyBuffer(targetConn, sourceConn, buf)
			receivedBytesCount += bytesCopied
			if err == nil {
				logrus.Debugf("Connection closed '%v' -> '%v' cid '%v', bytes transferred '%v' (%v).", sourceConn, targetConn, cid, bytesCopied, receivedBytesCount)
			} else {
				logrus.Debugf("Connection closed '%v' -> '%v' cid '%v', bytes transferred '%v' (%v), error: %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, bytesCopied, receivedBytesCount, err)
			}
			return
		}

		if !keepAliveMode {
			return
		}
	}

}

func startProxyTCP(cid ConnectionID, targetConn net.Conn, sourceConn net.Conn) {
	logrus.Infof("Start tcp-proxy connection from '%v' to'%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)

	go func() {
		buf := netbufGet()
		defer netbufPut(buf)

		_, err := io.CopyBuffer(targetConn, sourceConn, buf)
		logrus.Debugf("Connection closed with error2 '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, err)
		sourceConn.Close()
		targetConn.Close()
	}()
	go func() {
		buf := netbufGet()
		defer netbufPut(buf)

		_, err := io.CopyBuffer(sourceConn, targetConn, buf)
		logrus.Debugf("Connection closed with error3 '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, err)
		sourceConn.Close()
		targetConn.Close()
	}()
}
<commit_msg>change message "Can't local lookup ip for domain " level<commit_after>package main

import (
	"bytes"
	"github.com/Sirupsen/logrus"
	"io"
	"net"
	"strconv"
	"sync"
)

const (
	NETBUF_SIZE = 2048 // bytes
)

//go:generate stringer -type KeepAliveModeType
type KeepAliveModeType int

const (
	KEEPALIVE_BOTH KeepAliveModeType = iota
	KEEPALIVE_NOBACKEND
)

const (
	KEEPALIVE_BOTH_STRING      = "both"
	KEEPALIVE_NOBACKEND_STRING = "nobackend"
)

var (
	poolNetBuffers sync.Pool
	keepAliveMode  KeepAliveModeType
)

// var-constants
var (
	HEAD_CONNECTION            = []byte("CONNECTION")
	HEAD_CONNECTION_KEEP_ALIVE = []byte("KEEP-ALIVE")
	HEAD_CONTENT_LENGTH        = []byte("CONTENT-LENGTH")
)

// Get or create network buffer for proxy
func netbufGet() (buf []byte) {
	bufInterface := poolNetBuffers.Get()
	if bufInterface == nil {
		buf = make([]byte, NETBUF_SIZE)
	} else {
		buf = bufInterface.([]byte)
		// prevent data leak
		for i := range buf {
			buf[i] = 0
		}

	}
	return buf
}

func netbufPut(buf []byte) {
	poolNetBuffers.Put(buf)
}

func proxyHTTPHeaders(cid ConnectionID, targetConn net.Conn, sourceConn net.Conn) (
	res struct {
		KeepAlive        bool
		ContentLength    int64
		HasContextLength bool
		Err              error
	}) {
	buf := netbufGet()
	defer netbufPut(buf)
	var totalReadBytes int

	// Read lines
readHeaderLines:
	for {
		var i int
		var headerStart []byte
		for i = 0; i < len(buf); i++ {
			var readBytes int
			readBytes, res.Err = sourceConn.Read(buf[i : i+1])
			totalReadBytes += readBytes
			if res.Err != nil {
				logrus.Debugf("Error while read header from '%v' cid '%v': %v", sourceConn.RemoteAddr(), cid, res.Err)
				return
			}
			if readBytes != 1 {
				logrus.Infof("Can't read a byte from header from '%v' cid '%v'", sourceConn.RemoteAddr(), cid)
				return
			}
			if buf[i] == ':' || buf[i] == '\n' {
				headerStart = buf[:i+1]
				logrus.Debugf("Header Name '%v' -> '%v' cid '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, buf[:i])
				break
			}
		}
		if len(headerStart) == 0 {
			logrus.Infof("Header line longer then buffer (%v bytes). Force close connection. '%v' -> '%v' cid '%v'.", len(buf), sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)
			return
		}

		// Empty line - end http headers
		if bytes.Equal(headerStart, []byte("\n")) || bytes.Equal(headerStart, []byte("\r\n")) {
			break readHeaderLines
		}

		headerName := headerStart[:len(headerStart)-1] // Cut trailing colon from start
		headerNameUpperCase := bytes.ToUpper(headerName)

		skipHeader := false
		for _, ownHeader := range cutHeaders {
			if bytes.Equal(ownHeader, headerNameUpperCase) {
				skipHeader = true
				break
			}
			ownHeaderS := string(ownHeader)
			headerNameS := string(headerNameUpperCase)
			skipHeader = ownHeaderS == headerNameS
		}

		if skipHeader {
			logrus.Debugf("Skip header: '%v' -> '%v' cid '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerNameUpperCase)
			buf[0] = headerStart[len(headerStart)-1]

			for buf[0] != '\n' {
				_, res.Err = sourceConn.Read(buf[:1])
				if res.Err != nil {
					logrus.Infof("Error read header. Close connections. '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, res.Err)
					sourceConn.Close()
					targetConn.Close()
					return
				}
			}
			continue readHeaderLines
		}

		logrus.Debugf("Copy header: '%v' -> '%v' cid '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerName)

		// copy header without changes
		_, res.Err = targetConn.Write(headerStart)
		if res.Err != nil {
			logrus.Infof("Write header start, from '%v' to '%v' cid '%v', headerStart='%s': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerStart, res.Err)
			return
		}

		needHeaderContent := bytes.Equal(headerNameUpperCase, HEAD_CONTENT_LENGTH) || bytes.Equal(headerNameUpperCase, HEAD_CONNECTION)
		headerContent := bytes.NewBuffer(buf[1:])
		headerContent.Reset()

		for buf[0] != '\n' {
			var readBytes int
			readBytes, res.Err = sourceConn.Read(buf[:1])
			if res.Err != nil {
				logrus.Infof("Error read header to copy. Close connections. '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, res.Err)
				return
			}
			if readBytes != 1 {
				logrus.Infof("Header copy read bytes != 1. Error. Close connections. '%v' -> '%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)
				return
			}
			_, res.Err = targetConn.Write(buf[:1])
			if res.Err != nil {
				logrus.Infof("Error write header. Close connections. '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, res.Err)
				return
			}
			if needHeaderContent {
				headerContent.WriteByte(buf[0])
			}
		}
		if needHeaderContent {
			switch {
			case bytes.Equal(headerNameUpperCase, HEAD_CONNECTION):
				res.KeepAlive = bytes.EqualFold(HEAD_CONNECTION_KEEP_ALIVE, bytes.TrimSpace(headerContent.Bytes()))

			case bytes.Equal(headerNameUpperCase, HEAD_CONTENT_LENGTH):
				res.ContentLength, res.Err = strconv.ParseInt(string(bytes.TrimSpace(headerContent.Bytes())), 10, 64)
				if res.Err == nil {
					res.HasContextLength = true
					logrus.Debugf("Header content-length parsed from '%v' to '%v' cid '%v': %v", sourceConn.RemoteAddr(),
						targetConn.RemoteAddr(), cid, res.ContentLength)
				} else {
					logrus.Infof("Can't header content-length parsed from '%v' to '%v' cid '%v' content '%s': %v", sourceConn.RemoteAddr(),
						targetConn.RemoteAddr(), cid, headerContent.Bytes(), res.Err)
				}

			default:
				logrus.Debugf("ERROR. Unknow why i need header content. Code error. From '%v' to '%v' cid '%v', header name '%s'",
					sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, headerNameUpperCase,
				)
			}
		}
	}

	remoteTcpAddr := sourceConn.RemoteAddr().(*net.TCPAddr)
	remoteAddrString := remoteTcpAddr.IP.String()

	headerBuf := bytes.NewBuffer(buf)
	headerBuf.Reset()

	// Write real IP
	for _, header := range realIPHeaderNames {
		headerBuf.Write(header)
		headerBuf.WriteByte(':')
		headerBuf.WriteString(remoteAddrString)
		headerBuf.WriteString("\r\n")
	}

	// Write CID
	if *connectionIdHeader != "" {
		headerBuf.WriteString(*connectionIdHeader)
		headerBuf.WriteString(": ")
		headerBuf.WriteString(cid.String())
		headerBuf.WriteString("\r\n")
	}

	headerBuf.Write(additionalHeaders)
	headerBuf.Write([]byte("\r\n")) // end http headers
	logrus.Debugf("Add headers. '%v' -> '%v': '%s'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), headerBuf.Bytes())

	_, res.Err = targetConn.Write(headerBuf.Bytes())
	if res.Err != nil {
		logrus.Infof("Error while write real ip headers to target '%v' -> '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), res.Err)
	}

	return
}

func startProxy(cid ConnectionID, targetAddr net.TCPAddr, in net.Conn) {
	targetConnCommon, err := net.DialTimeout("tcp", targetAddr.String(), *targetConnTimeout)
	if err != nil {
		logrus.Warnf("Can't connect to target '%v' cid '%v': %v", targetAddr.String(), cid, err)
		in.Close()
		return
	}

	targetConn, ok := targetConnCommon.(*net.TCPConn)
	targetConn.SetKeepAlive(true)
	targetConn.SetKeepAlivePeriod(*tcpKeepAliveInterval)

	if !ok {
		logrus.Errorf("Can't cast connection to tcp connection, target '%v' cid '%v'", targetAddr.String(), cid)
		return
	}

	switch *proxyMode {
	case "http":
		startProxyHTTP(cid, targetConn, in)
	case "tcp":
		startProxyTCP(cid, targetConn, in)
	default:
		logrus.Panicf("Unknow proxy mode cid '%v': %v", cid, *proxyMode)
	}
}

func startProxyHTTP(cid ConnectionID, targetConn net.Conn, sourceConn net.Conn) {
	defer targetConn.Close()
	defer sourceConn.Close()

	logrus.Debugf("Start http-proxy connection from '%v' to '%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)

	buf := netbufGet()
	defer netbufPut(buf)
	var receivedBytesCount int64
	for {
		state := proxyHTTPHeaders(cid, targetConn, sourceConn)
		keepAlive := state.KeepAlive
		if state.Err != nil {
			logrus.Debug("Cid '%v'. Can't read headers: %v", cid, state.Err)
			return
		}
		if state.HasContextLength {
			logrus.Debugf("Start keep-alieved proxy. '%v' -> '%v' cid '%v', content-length '%v'", sourceConn.RemoteAddr(),
				targetConn.RemoteAddr(), cid, state.ContentLength)

			//request
			bytesCopied, err := io.CopyBuffer(targetConn, io.LimitReader(sourceConn, state.ContentLength), buf)
			receivedBytesCount += bytesCopied

			if err == nil {
				logrus.Debugf("Connection chunk copied '%v' -> '%v' cid '%v', bytes transferred '%v' (%v), error: %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, bytesCopied, receivedBytesCount, err)
			} else {
				logrus.Debugf("Connection closed '%v' -> '%v' cid '%v', bytes transferred '%v' (%v), error: %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, bytesCopied, receivedBytesCount, err)
				return
			}

			// proxy answer
			answerState := proxyHTTPHeaders(cid, sourceConn, targetConn)
			keepAlive = keepAlive && answerState.KeepAlive
			var reader io.Reader

			if answerState.HasContextLength {
				reader = io.LimitReader(targetConn, answerState.ContentLength)
			} else {
				reader = targetConn
			}

			_, err = io.CopyBuffer(sourceConn, reader, buf)
			if err != nil {
				logrus.Debugf("Cid '%v'. Can't copy answer to customer: %v", cid, err)
				return
			}
		} else {
			// answer from server proxy without changes
			go func() {
				buf := netbufGet()
				defer netbufPut(buf)

				_, err := io.CopyBuffer(sourceConn, targetConn, buf)
				logrus.Debugf("Connection closed with error1 '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, err)
				sourceConn.Close()
				targetConn.Close()
			}()

			logrus.Debugf("Start proxy without support keepalive middle headers '%v' -> '%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)
			bytesCopied, err := io.CopyBuffer(targetConn, sourceConn, buf)
			receivedBytesCount += bytesCopied
			if err == nil {
				logrus.Debugf("Connection closed '%v' -> '%v' cid '%v', bytes transferred '%v' (%v).", sourceConn, targetConn, cid, bytesCopied, receivedBytesCount)
			} else {
				logrus.Debugf("Connection closed '%v' -> '%v' cid '%v', bytes transferred '%v' (%v), error: %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, bytesCopied, receivedBytesCount, err)
			}
			return
		}

		if !keepAlive {
			return
		}
	}

}

func startProxyTCP(cid ConnectionID, targetConn net.Conn, sourceConn net.Conn) {
	logrus.Infof("Start tcp-proxy connection from '%v' to'%v' cid '%v'", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid)

	go func() {
		buf := netbufGet()
		defer netbufPut(buf)

		_, err := io.CopyBuffer(targetConn, sourceConn, buf)
		logrus.Debugf("Connection closed with error2 '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, err)
		sourceConn.Close()
		targetConn.Close()
	}()
	go func() {
		buf := netbufGet()
		defer netbufPut(buf)

		_, err := io.CopyBuffer(sourceConn, targetConn, buf)
		logrus.Debugf("Connection closed with error3 '%v' -> '%v' cid '%v': %v", sourceConn.RemoteAddr(), targetConn.RemoteAddr(), cid, err)
		sourceConn.Close()
		targetConn.Close()
	}()
}
<|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")

	replaceContextPool = sync.Pool{New: func() interface{} {
		return &replacer.ReplaceContext{}
	}}
)

type ProxyConfig struct {
	ReduceMemoryUsage      bool
	BaseDomain             string
	BaseStaticDomain       string
	LogVerbosity           int
	GzipUpstream           bool
	FilterFeed             bool
	AddUselessProxyMessage 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},
			ReadTimeout:               30 * time.Second,
			WriteTimeout:              10 * time.Second,
			DisablePathNormalizing:    true,
			NoDefaultUserAgentHeader:  true,
			MaxIdemponentCallAttempts: 0,
			RetryIf: func(request *fasthttp.Request) bool {
				return false
			},
		},
		replacer: &replacer.Replacer{
			ProxyBaseDomain:        config.BaseDomain,
			ProxyStaticDomain:      config.BaseStaticDomain,
			FilterFeed:             config.FilterFeed,
			AddUselessProxyMessage: config.AddUselessProxyMessage,
		},
		tracker: &tracker{
			uniqueUsers: make(map[string]bool),
		},
		config: config,
	}
	p.server = &fasthttp.Server{
		Handler:                      p.handleProxy,
		ReduceMemoryUsage:            config.ReduceMemoryUsage,
		ReadBufferSize:               readBufferSize,
		ReadTimeout:                  10 * time.Second,
		WriteTimeout:                 20 * time.Second,
		IdleTimeout:                  1 * time.Minute,
		NoDefaultContentType:         true,
		DisablePreParseMultipartForm: true,
		Name:                         "vk-proxy",
	}
	p.tracker.server = p.server
	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 := replaceContextPool.Get().(*replacer.ReplaceContext)
	replaceContext.RequestCtx = ctx
	replaceContext.Method = ctx.Method()
	replaceContext.OriginHost = string(ctx.Request.Host())

	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.Do(&ctx.Request, &ctx.Response)
	if err == nil {
		err = p.processProxyResponse(ctx, replaceContext)
	}

	replaceContext.Reset()
	replaceContextPool.Put(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(fasthttp.HeaderXForwardedFor) // 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") &&
			!strings.HasSuffix(endpoint, ".mycdn.me"){
			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.SetBytesV(fasthttp.HeaderAcceptEncoding, gzip)
	} else {
		req.Header.Del(fasthttp.HeaderAcceptEncoding)
	}

	req.Header.Del(fasthttp.HeaderConnection)
	return true
}

func (p *Proxy) processProxyResponse(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) error {
	res := &ctx.Response
	res.Header.Del(fasthttp.HeaderSetCookie)
	res.Header.Del(fasthttp.HeaderConnection)
	res.Header.SetBytesV(fasthttp.HeaderServer, vkProxyName)

	var buf *bytebufferpool.ByteBuffer
	// Gunzip body if needed
	if bytes.Contains(res.Header.Peek(fasthttp.HeaderContentEncoding), gzip) {
		res.Header.Del(fasthttp.HeaderContentEncoding)
		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
	server      *fasthttp.Server
}

func (t *tracker) start() {
	go func() {
		for range time.Tick(60 * time.Second) {
			t.lock.Lock()
			log.Printf("Requests: %d, Traffic: %s, Online: %d, Concurrency: %d",
				t.requests, bytefmt.ByteSize(t.bytes), len(t.uniqueUsers),
				t.server.GetCurrentConcurrency(),
			)
			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>Добавил api.ok.ru в белый список прокси для работы звонков 2.0<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")

	replaceContextPool = sync.Pool{New: func() interface{} {
		return &replacer.ReplaceContext{}
	}}
)

type ProxyConfig struct {
	ReduceMemoryUsage      bool
	BaseDomain             string
	BaseStaticDomain       string
	LogVerbosity           int
	GzipUpstream           bool
	FilterFeed             bool
	AddUselessProxyMessage 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},
			ReadTimeout:               30 * time.Second,
			WriteTimeout:              10 * time.Second,
			DisablePathNormalizing:    true,
			NoDefaultUserAgentHeader:  true,
			MaxIdemponentCallAttempts: 0,
			RetryIf: func(request *fasthttp.Request) bool {
				return false
			},
		},
		replacer: &replacer.Replacer{
			ProxyBaseDomain:        config.BaseDomain,
			ProxyStaticDomain:      config.BaseStaticDomain,
			FilterFeed:             config.FilterFeed,
			AddUselessProxyMessage: config.AddUselessProxyMessage,
		},
		tracker: &tracker{
			uniqueUsers: make(map[string]bool),
		},
		config: config,
	}
	p.server = &fasthttp.Server{
		Handler:                      p.handleProxy,
		ReduceMemoryUsage:            config.ReduceMemoryUsage,
		ReadBufferSize:               readBufferSize,
		ReadTimeout:                  10 * time.Second,
		WriteTimeout:                 20 * time.Second,
		IdleTimeout:                  1 * time.Minute,
		NoDefaultContentType:         true,
		DisablePreParseMultipartForm: true,
		Name:                         "vk-proxy",
	}
	p.tracker.server = p.server
	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 := replaceContextPool.Get().(*replacer.ReplaceContext)
	replaceContext.RequestCtx = ctx
	replaceContext.Method = ctx.Method()
	replaceContext.OriginHost = string(ctx.Request.Host())

	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.Do(&ctx.Request, &ctx.Response)
	if err == nil {
		err = p.processProxyResponse(ctx, replaceContext)
	}

	replaceContext.Reset()
	replaceContextPool.Put(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(fasthttp.HeaderXForwardedFor) // 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") &&
			!strings.HasSuffix(endpoint, ".mycdn.me") &&
			endpoint != "api.ok.ru" {
			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.SetBytesV(fasthttp.HeaderAcceptEncoding, gzip)
	} else {
		req.Header.Del(fasthttp.HeaderAcceptEncoding)
	}

	req.Header.Del(fasthttp.HeaderConnection)
	return true
}

func (p *Proxy) processProxyResponse(ctx *fasthttp.RequestCtx, replaceContext *replacer.ReplaceContext) error {
	res := &ctx.Response
	res.Header.Del(fasthttp.HeaderSetCookie)
	res.Header.Del(fasthttp.HeaderConnection)
	res.Header.SetBytesV(fasthttp.HeaderServer, vkProxyName)

	var buf *bytebufferpool.ByteBuffer
	// Gunzip body if needed
	if bytes.Contains(res.Header.Peek(fasthttp.HeaderContentEncoding), gzip) {
		res.Header.Del(fasthttp.HeaderContentEncoding)
		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
	server      *fasthttp.Server
}

func (t *tracker) start() {
	go func() {
		for range time.Tick(60 * time.Second) {
			t.lock.Lock()
			log.Printf("Requests: %d, Traffic: %s, Online: %d, Concurrency: %d",
				t.requests, bytefmt.ByteSize(t.bytes), len(t.uniqueUsers),
				t.server.GetCurrentConcurrency(),
			)
			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|>