<commit_before>package main

import (
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"path"
	"strconv"

	"github.com/bruston/handlers/gzipped"
)

const (
	errWantInteger           = "n must be an integer"
	errStreamingNotSupported = "your client does not support streaming"
	maxBytes                 = 102400
	maxLines                 = 100
)

func defaultHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if o := r.Header.Get("Origin"); o != "" {
			w.Header().Set("Access-Control-Allow-Origin", o)
			w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
			w.Header().Set("Access-Control-Allow-Headers",
				"Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")
			if r.Method == "OPTIONS" {
				return
			}
		}
		h.ServeHTTP(w, r)
	})
}

func main() {
	listen := flag.String("listen", "127.0.0.1:8090", "The host and port to listen on.")
	flag.Parse()
	http.HandleFunc("/headers", headers)
	http.HandleFunc("/status/", status)
	http.HandleFunc("/ip", ip)
	http.HandleFunc("/get", get)
	http.Handle("/gzip", gzipped.New(http.HandlerFunc(gzip)))
	http.HandleFunc("/user-agent", userAgent)
	http.HandleFunc("/bytes/", writeBytes)
	http.HandleFunc("/stream/", stream)
	log.Fatal(http.ListenAndServe(*listen, defaultHandler(http.DefaultServeMux)))
}

func jsonHeader(w http.ResponseWriter) {
	w.Header().Set("Content-type", "application/json")
}

func writeJSON(w http.ResponseWriter, data interface{}, code int) error {
	jsonHeader(w)
	w.WriteHeader(code)
	return json.NewEncoder(w).Encode(data)
}

func headers(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, r.Header, http.StatusOK)
}

func status(w http.ResponseWriter, r *http.Request) {
	code, err := strconv.Atoi(path.Base(r.URL.Path))
	if err != nil {
		http.Error(w, "status code must be an integer", http.StatusBadRequest)
		return
	}
	w.WriteHeader(code)
}

type origin struct {
	IP           string `json:"ip"`
	ForwardedFor string `json:"forwarded_for,omitempty"`
}

func getOrigin(r *http.Request) origin {
	host, _, _ := net.SplitHostPort(r.RemoteAddr)
	return origin{host, r.Header.Get("X-Forwarded-For")}
}

func ip(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, getOrigin(r), http.StatusOK)
}

type request struct {
	Args    url.Values  `json:"args"`
	Gzipped bool        `json:"gzipped,omitempty"`
	Headers http.Header `json:"headers"`
	Origin  origin      `json:"origin"`
	URL     string      `json:"url"`
}

func rawURL(r *http.Request) string {
	var scheme string
	if r.TLS == nil {
		scheme = "http"
	} else {
		scheme = "https"
	}
	return scheme + "://" + r.Host + r.URL.String()
}

func getReq(r *http.Request) request {
	return request{
		Args:    r.URL.Query(),
		Headers: r.Header,
		Origin:  getOrigin(r),
		URL:     rawURL(r),
	}
}

func get(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
		return
	}
	writeJSON(w, getReq(r), http.StatusOK)
}

func gzip(w http.ResponseWriter, r *http.Request) {
	req := getReq(r)
	if _, ok := w.(gzipped.GzipResponseWriter); ok {
		req.Gzipped = true
	}
	writeJSON(w, req, http.StatusOK)
}

func userAgent(w http.ResponseWriter, r *http.Request) {
	var resp struct {
		UserAgent string `json:"user-agent"`
	}
	resp.UserAgent = r.Header.Get("User-Agent")
	writeJSON(w, resp, http.StatusOK)
}

func writeBytes(w http.ResponseWriter, r *http.Request) {
	n, err := strconv.Atoi(path.Base(r.URL.Path))
	if err != nil || n < 0 || n > maxBytes {
		http.Error(w, fmt.Sprintf("number of bytes must be in range: 0 - %d", maxBytes), http.StatusBadRequest)
		return
	}
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write(b)
}

func min(a, b int) int {
	if a <= b {
		return a
	}
	return b
}

func stream(w http.ResponseWriter, r *http.Request) {
	n, err := strconv.Atoi(path.Base(r.URL.Path))
	if err != nil || n < 0 {
		http.Error(w, errWantInteger, http.StatusBadRequest)
		return
	}
	n = min(n, maxLines)
	f, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, errStreamingNotSupported, http.StatusBadRequest)
		return
	}
	req := getReq(r)
	jsonHeader(w)
	for i := 0; i < n; i++ {
		if err := json.NewEncoder(w).Encode(req); err != nil {
			return
		}
		f.Flush()
	}
}
<commit_msg>Simplify getOrigin<commit_after>package main

import (
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"path"
	"strconv"

	"github.com/bruston/handlers/gzipped"
)

const (
	errWantInteger           = "n must be an integer"
	errStreamingNotSupported = "your client does not support streaming"
	maxBytes                 = 102400
	maxLines                 = 100
)

func defaultHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if o := r.Header.Get("Origin"); o != "" {
			w.Header().Set("Access-Control-Allow-Origin", o)
			w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
			w.Header().Set("Access-Control-Allow-Headers",
				"Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")
			if r.Method == "OPTIONS" {
				return
			}
		}
		h.ServeHTTP(w, r)
	})
}

func main() {
	listen := flag.String("listen", "127.0.0.1:8090", "The host and port to listen on.")
	flag.Parse()
	http.HandleFunc("/headers", headers)
	http.HandleFunc("/status/", status)
	http.HandleFunc("/ip", ip)
	http.HandleFunc("/get", get)
	http.Handle("/gzip", gzipped.New(http.HandlerFunc(gzip)))
	http.HandleFunc("/user-agent", userAgent)
	http.HandleFunc("/bytes/", writeBytes)
	http.HandleFunc("/stream/", stream)
	log.Fatal(http.ListenAndServe(*listen, defaultHandler(http.DefaultServeMux)))
}

func jsonHeader(w http.ResponseWriter) {
	w.Header().Set("Content-type", "application/json")
}

func writeJSON(w http.ResponseWriter, data interface{}, code int) error {
	jsonHeader(w)
	w.WriteHeader(code)
	return json.NewEncoder(w).Encode(data)
}

func headers(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, r.Header, http.StatusOK)
}

func status(w http.ResponseWriter, r *http.Request) {
	code, err := strconv.Atoi(path.Base(r.URL.Path))
	if err != nil {
		http.Error(w, "status code must be an integer", http.StatusBadRequest)
		return
	}
	w.WriteHeader(code)
}

func getOrigin(r *http.Request) string {
	host, _, _ := net.SplitHostPort(r.RemoteAddr)
	if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" && forwarded != host {
		host = fmt.Sprintf("%s, %s", forwarded, host)
	}
	return host
}

func ip(w http.ResponseWriter, r *http.Request) {
	var o struct {
		Origin string `json:"origin"`
	}
	o.Origin = getOrigin(r)
	writeJSON(w, o, http.StatusOK)
}

type request struct {
	Args    url.Values  `json:"args"`
	Gzipped bool        `json:"gzipped,omitempty"`
	Headers http.Header `json:"headers"`
	Origin  string      `json:"origin"`
	URL     string      `json:"url"`
}

func rawURL(r *http.Request) string {
	var scheme string
	if r.TLS == nil {
		scheme = "http"
	} else {
		scheme = "https"
	}
	return scheme + "://" + r.Host + r.URL.String()
}

func getReq(r *http.Request) request {
	return request{
		Args:    r.URL.Query(),
		Headers: r.Header,
		Origin:  getOrigin(r),
		URL:     rawURL(r),
	}
}

func get(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
		return
	}
	writeJSON(w, getReq(r), http.StatusOK)
}

func gzip(w http.ResponseWriter, r *http.Request) {
	req := getReq(r)
	if _, ok := w.(gzipped.GzipResponseWriter); ok {
		req.Gzipped = true
	}
	writeJSON(w, req, http.StatusOK)
}

func userAgent(w http.ResponseWriter, r *http.Request) {
	var resp struct {
		UserAgent string `json:"user-agent"`
	}
	resp.UserAgent = r.Header.Get("User-Agent")
	writeJSON(w, resp, http.StatusOK)
}

func writeBytes(w http.ResponseWriter, r *http.Request) {
	n, err := strconv.Atoi(path.Base(r.URL.Path))
	if err != nil || n < 0 || n > maxBytes {
		http.Error(w, fmt.Sprintf("number of bytes must be in range: 0 - %d", maxBytes), http.StatusBadRequest)
		return
	}
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write(b)
}

func min(a, b int) int {
	if a <= b {
		return a
	}
	return b
}

func stream(w http.ResponseWriter, r *http.Request) {
	n, err := strconv.Atoi(path.Base(r.URL.Path))
	if err != nil || n < 0 {
		http.Error(w, errWantInteger, http.StatusBadRequest)
		return
	}
	n = min(n, maxLines)
	f, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, errStreamingNotSupported, http.StatusBadRequest)
		return
	}
	req := getReq(r)
	jsonHeader(w)
	for i := 0; i < n; i++ {
		if err := json.NewEncoder(w).Encode(req); err != nil {
			return
		}
		f.Flush()
	}
}
<|endoftext|>


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

   This file is part of Httpgzip.

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

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

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

// Package httpgzip implements an http.Handler which wraps an existing
// http.Handler adding Gzip compression for appropriate requests.
//
// It attempts to properly parse the request's Accept-Encoding header
// according to RFC 2616 and does not just do a
// strings.Contains(header,"gzip"). It will serve either gzip or
// identity content codings (or return 406 Not Acceptable status if
// can do neither).
//
// It works correctly with handlers such as http.FileServer which
// honour Range requests by removing the Range header when requests
// prefer gzip encoding. This is necessary since Range applies to the
// Gzipped content and the wrapped handler is not aware of the
// compression when it writes byte ranges.
package httpgzip // import "xi2.org/x/httpgzip"

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

// ErrStatusNotAcceptable is returned by Write calls to the
// http.ResponseWriter within the wrapped http.Handler if a
// http.StatusNotAcceptable has been written by the outer Handler.
var ErrStatusNotAcceptable = errors.New("httpgzip: http.StatusNotAcceptable has been written")

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

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

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

// A gzipResponseWriter is a modified http.ResponseWriter. If the
// content to be written is of a type contained in its contentTypes
// field and the request allows and prefers Gzip compression then the
// response is compressed and the Content-Encoding header is
// set. Otherwise a gzipResponseWriter behaves mostly like a normal
// http.ResponseWriter. It is important to call the Close method when
// writing is finished in order to flush and close the Writer.
type gzipResponseWriter struct {
	http.ResponseWriter
	httpStatus   int
	contentTypes map[string]struct{}
	encs         []encoding
	gw           *gzip.Writer
	buf          *bytes.Buffer
}

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

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

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

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

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

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

const (
	encIdentity encoding = iota
	encGzip
)

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

// acceptIdentity returns true if identity encoding is accepted.
func acceptIdentity(encs []encoding) bool {
	for _, e := range encs {
		if e == encIdentity {
			return true
		}
	}
	return false
}

// preferGzipToIdentity returns true if gzip encoding is accepted and
// preferred to identity encoding.
func preferGzipToIdentity(encs []encoding) bool {
	return len(encs) > 0 && encs[0] == encGzip
}

// NewHandler returns a new http.Handler which wraps a handler h
// adding Gzip compression to responses whose content types are in
// contentTypes (unless the corresponding request does not allow or
// prefer Gzip compression). If contentTypes is nil then it is set to
// DefaultContentTypes.
//
// The new http.Handler sets the Content-Encoding, Vary and
// Content-Type headers in its responses as appropriate. If the
// request expresses a preference for gzip encoding then any "Range"
// headers are removed from the request before forwarding it to
// h. This happens regardless of whether gzip encoding is eventually
// used in the response or not.
func NewHandler(h http.Handler, contentTypes map[string]struct{}) http.Handler {
	if contentTypes == nil {
		contentTypes = DefaultContentTypes
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// add Vary header
		w.Header().Add("Vary", "Accept-Encoding")
		// check client's accepted encodings
		encs := acceptedEncodings(r)
		if preferGzipToIdentity(encs) {
			// cannot accept Range requests for possibly gzipped
			// responses
			r.Header.Del("Range")
		}
		w = newGzipResponseWriter(w, contentTypes, encs)
		defer w.(*gzipResponseWriter).Close()
		// call original handler's ServeHTTP
		h.ServeHTTP(w, r)
	})
}
<commit_msg>Do not return ErrStatusNotAcceptable in Close method<commit_after>/*
   Copyright 2015 The Httpgzip Authors. See the AUTHORS file at the
   top-level directory of this distribution and at
   <https://xi2.org/x/httpgzip/m/AUTHORS>.

   This file is part of Httpgzip.

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

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

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

// Package httpgzip implements an http.Handler which wraps an existing
// http.Handler adding Gzip compression for appropriate requests.
//
// It attempts to properly parse the request's Accept-Encoding header
// according to RFC 2616 and does not just do a
// strings.Contains(header,"gzip"). It will serve either gzip or
// identity content codings (or return 406 Not Acceptable status if
// can do neither).
//
// It works correctly with handlers such as http.FileServer which
// honour Range requests by removing the Range header when requests
// prefer gzip encoding. This is necessary since Range applies to the
// Gzipped content and the wrapped handler is not aware of the
// compression when it writes byte ranges.
package httpgzip // import "xi2.org/x/httpgzip"

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

// ErrStatusNotAcceptable is returned by Write calls to the
// http.ResponseWriter within the wrapped http.Handler if a
// http.StatusNotAcceptable has been written by the outer Handler.
var ErrStatusNotAcceptable = errors.New("httpgzip: http.StatusNotAcceptable has been written")

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

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

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

// A gzipResponseWriter is a modified http.ResponseWriter. If the
// content to be written is of a type contained in its contentTypes
// field and the request allows and prefers Gzip compression then the
// response is compressed and the Content-Encoding header is
// set. Otherwise a gzipResponseWriter behaves mostly like a normal
// http.ResponseWriter. It is important to call the Close method when
// writing is finished in order to flush and close the Writer.
type gzipResponseWriter struct {
	http.ResponseWriter
	httpStatus   int
	contentTypes map[string]struct{}
	encs         []encoding
	gw           *gzip.Writer
	buf          *bytes.Buffer
}

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

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

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

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

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

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

const (
	encIdentity encoding = iota
	encGzip
)

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

// acceptIdentity returns true if identity encoding is accepted.
func acceptIdentity(encs []encoding) bool {
	for _, e := range encs {
		if e == encIdentity {
			return true
		}
	}
	return false
}

// preferGzipToIdentity returns true if gzip encoding is accepted and
// preferred to identity encoding.
func preferGzipToIdentity(encs []encoding) bool {
	return len(encs) > 0 && encs[0] == encGzip
}

// NewHandler returns a new http.Handler which wraps a handler h
// adding Gzip compression to responses whose content types are in
// contentTypes (unless the corresponding request does not allow or
// prefer Gzip compression). If contentTypes is nil then it is set to
// DefaultContentTypes.
//
// The new http.Handler sets the Content-Encoding, Vary and
// Content-Type headers in its responses as appropriate. If the
// request expresses a preference for gzip encoding then any "Range"
// headers are removed from the request before forwarding it to
// h. This happens regardless of whether gzip encoding is eventually
// used in the response or not.
func NewHandler(h http.Handler, contentTypes map[string]struct{}) http.Handler {
	if contentTypes == nil {
		contentTypes = DefaultContentTypes
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// add Vary header
		w.Header().Add("Vary", "Accept-Encoding")
		// check client's accepted encodings
		encs := acceptedEncodings(r)
		if preferGzipToIdentity(encs) {
			// cannot accept Range requests for possibly gzipped
			// responses
			r.Header.Del("Range")
		}
		w = newGzipResponseWriter(w, contentTypes, encs)
		defer w.(*gzipResponseWriter).Close()
		// call original handler's ServeHTTP
		h.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	flag.StringVar(&env, "e", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})
}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	fmt.Println("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	destIndex := rand.Intn(len(route.Destinations))
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		fmt.Println("error!", err)
		panic(err)
	}
	fmt.Println("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		fmt.Println("Error proxying!", err)
		etype := reflect.TypeOf(err)
		fmt.Println("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 1 {
				fmt.Println("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
				return
			} else {
				fmt.Println("It's a network error and no other destinations available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
			}
			// start new worker
			payload := map[string]interface{}{
				"token":      route.Token,
				"project_id": route.ProjectId,
				"code_name":  route.CodeName,
			}
			workerapi := worker.New()
			workerapi.Settings.UseConfigMap(payload)
			jsonPayload, err := json.Marshal(payload)
			if err != nil {
				fmt.Println("Couldn't marshal json!", err)
				return
			}
			timeout := time.Second*120
			task := worker.Task{
				CodeName: route.CodeName,
				Payload:  string(jsonPayload),
				Timeout:  &timeout, // let's have these die quickly while testing
			}
			tasks := make([]worker.Task, 1)
			tasks[0] = task
			taskIds, err := workerapi.TaskQueue(tasks...)
			fmt.Println("Tasks queued.", taskIds)
			if err != nil {
				fmt.Println("Couldn't queue up worker!", err)
				return
			}
		}
		// start new worker if it's a connection error
		return
	}
	fmt.Println("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	fmt.Println("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		fmt.Println("registered route:", route)

		putRoute(route)

	} else {
		r2 := Route2{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&r2)
		// todo: do we need to close body?
		fmt.Println("DECODED:", r2)

		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln(w, "This host is not registered!", err))
			return
			//			route = &Route{}
		}
		fmt.Println("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		fmt.Println("ROUTE new:", route)

		putRoute(route)
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (Route, error) {
	rx, err := icache.Get(host)
	route := Route{}
	if err == nil {
		route = rx.(Route)
	}
	return route, err
}

func putRoute(route Route) {
	item := cache.Item{}
	item.Value = route
	icache.Put(route.Host, &item)
}
<commit_msg>Moved stuff to main instead of init<commit_after>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	flag.StringVar(&env, "e", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	fmt.Println("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	destIndex := rand.Intn(len(route.Destinations))
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		fmt.Println("error!", err)
		panic(err)
	}
	fmt.Println("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		fmt.Println("Error proxying!", err)
		etype := reflect.TypeOf(err)
		fmt.Println("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 1 {
				fmt.Println("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
				return
			} else {
				fmt.Println("It's a network error and no other destinations available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
			}
			// start new worker
			payload := map[string]interface{}{
				"token":      route.Token,
				"project_id": route.ProjectId,
				"code_name":  route.CodeName,
			}
			workerapi := worker.New()
			workerapi.Settings.UseConfigMap(payload)
			jsonPayload, err := json.Marshal(payload)
			if err != nil {
				fmt.Println("Couldn't marshal json!", err)
				return
			}
			timeout := time.Second*120
			task := worker.Task{
				CodeName: route.CodeName,
				Payload:  string(jsonPayload),
				Timeout:  &timeout, // let's have these die quickly while testing
			}
			tasks := make([]worker.Task, 1)
			tasks[0] = task
			taskIds, err := workerapi.TaskQueue(tasks...)
			fmt.Println("Tasks queued.", taskIds)
			if err != nil {
				fmt.Println("Couldn't queue up worker!", err)
				return
			}
		}
		// start new worker if it's a connection error
		return
	}
	fmt.Println("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	fmt.Println("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		fmt.Println("registered route:", route)

		putRoute(route)

	} else {
		r2 := Route2{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&r2)
		// todo: do we need to close body?
		fmt.Println("DECODED:", r2)

		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln(w, "This host is not registered!", err))
			return
			//			route = &Route{}
		}
		fmt.Println("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		fmt.Println("ROUTE new:", route)

		putRoute(route)
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (Route, error) {
	rx, err := icache.Get(host)
	route := Route{}
	if err == nil {
		route = rx.(Route)
	}
	return route, err
}

func putRoute(route Route) {
	item := cache.Item{}
	item.Value = route
	icache.Put(route.Host, &item)
}
<|endoftext|>


<commit_before>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"github.com/iron-io/golog"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
//	"io/ioutil"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	// when this was e, it was erroring out.
	flag.StringVar(&env, "e", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)
	r.HandleFunc("/ping", Ping) // for health

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	golog.Infoln("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
 golog.Infoln("getting route for host:", host)
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln("Host not registered or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	dlen := len(route.Destinations)
	if dlen == 0 {
		golog.Infoln("No workers running, starting new task.")
		startNewWorker(route)
		common.SendError(w, 500, fmt.Sprintln("No workers running, starting them up..."))
		return
	}
	if dlen < 3 {
		golog.Infoln("Only one worker running, starting a new task.")
		startNewWorker(route)
	}
	destIndex := rand.Intn(dlen)
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		golog.Infoln("error!", err)
		panic(err)
	}
	golog.Infoln("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		golog.Infoln("Error proxying!", err)
		etype := reflect.TypeOf(err)
		golog.Infoln("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 3 { // always want at least two running
				golog.Infoln("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					golog.Infoln("couldn't update routing table 1", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table 1", err))
					return
				}
				fmt.Println("New route:", route)
				return
			} else {
				golog.Infoln("It's a network error and less than two other workers available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					golog.Infoln("couldn't update routing table:", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table", err))
					return
				}
				golog.Infoln("New route:", route)
			}
			// start new worker if it's a connection error
			startNewWorker(route)
		}
		return
	}
	golog.Infoln("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

func startNewWorker(route *Route) (error) {
	golog.Infoln("Starting a new worker")
	// start new worker
	payload := map[string]interface{}{
		"token":      route.Token,
		"project_id": route.ProjectId,
		"code_name":  route.CodeName,
	}
	workerapi := worker.New()
	workerapi.Settings.UseConfigMap(payload)
	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		golog.Infoln("Couldn't marshal json!", err)
		return err
	}
	timeout := time.Second * time.Duration(1800 + rand.Intn(600)) // a little random factor in here to spread out worker deaths
	task := worker.Task{
		CodeName: route.CodeName,
		Payload:  string(jsonPayload),
		Timeout:  &timeout, // let's have these die quickly while testing
	}
	tasks := make([]worker.Task, 1)
	tasks[0] = task
	taskIds, err := workerapi.TaskQueue(tasks...)
	golog.Infoln("Tasks queued.", taskIds)
	if err != nil {
		golog.Infoln("Couldn't queue up worker!", err)
		return err
	}
	return err
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

//	s, err := ioutil.ReadAll(req.Body)
//	fmt.Println("req.body:", err, string(s))

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	golog.Infoln("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		if !common.ReadJSON(w, req, &route) {
			return
		}
		golog.Infoln("body read into route:", route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		err := putRoute(&route)
		if err != nil {
			golog.Infoln("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		golog.Infoln("registered route:", route)
		fmt.Fprintln(w, "Host registered successfully.")

	} else {
		r2 := Route2{}
		if !common.ReadJSON(w, req, &r2) {
			return
		}
		// todo: do we need to close body?
		golog.Infoln("DECODED:", r2)
		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln("This host is not registered!", err))
			return
			//			route = &Route{}
		}
		golog.Infoln("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		golog.Infoln("ROUTE new:", route)
		err = putRoute(route)
		if err != nil {
			golog.Infoln("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (*Route, error) {
	rx, err := icache.Get(host)
	if err != nil {
		return nil, err
	}
	rx2 := []byte(rx.(string))
	route := Route{}
	err = json.Unmarshal(rx2, &route)
	if err != nil {
		return nil, err
	}
	return &route, err
}

func putRoute(route *Route) (error) {
	item := cache.Item{}
	v, err := json.Marshal(route)
	if err != nil {
		return err
	}
	item.Value = string(v)
	err = icache.Put(route.Host, &item)
	return err
}

func Ping(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintln(w, "pong")
}
<commit_msg>If can't connect to worker, retries with another endpoint.<commit_after>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"github.com/iron-io/golog"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
//	"io/ioutil"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	// when this was e, it was erroring out.
	flag.StringVar(&env, "e", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)
	r.HandleFunc("/ping", Ping) // for health

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	golog.Infoln("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
 	golog.Infoln("getting route for host:", host)
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln("Host not registered or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	dlen := len(route.Destinations)
	if dlen == 0 {
		golog.Infoln("No workers running, starting new task.")
		startNewWorker(route)
		common.SendError(w, 500, fmt.Sprintln("No workers running, starting them up..."))
		return
	}
	if dlen < 3 {
		golog.Infoln("Only one worker running, starting a new task.")
		startNewWorker(route)
	}
	serveEndpoint(w, req, route)
}

func serveEndpoint(w http.ResponseWriter, req *http.Request, route *Route) {
	dlen := len(route.Destinations)
	destIndex := rand.Intn(dlen)
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		golog.Infoln("error!", err)
		panic(err)
	}
	golog.Infoln("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		golog.Infoln("Error proxying!", err)
		etype := reflect.TypeOf(err)
		golog.Infoln("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 3 { // always want at least two running
				golog.Infoln("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					golog.Infoln("couldn't update routing table 1", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table 1", err))
					return
				}
				fmt.Println("New route:", route)
				// chooose another endpoint to serve
			} else {
				golog.Infoln("It's a network error and less than three other workers available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					golog.Infoln("couldn't update routing table:", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table", err))
					return
				}
				golog.Infoln("New route:", route)
				// start new worker if it's a connection error
				startNewWorker(route)
			}
			serveEndpoint(w, req, route)
			return
		}
		return
	}
	golog.Infoln("Served!")
}

func startNewWorker(route *Route) (error) {
	golog.Infoln("Starting a new worker")
	// start new worker
	payload := map[string]interface{}{
		"token":      route.Token,
		"project_id": route.ProjectId,
		"code_name":  route.CodeName,
	}
	workerapi := worker.New()
	workerapi.Settings.UseConfigMap(payload)
	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		golog.Infoln("Couldn't marshal json!", err)
		return err
	}
	timeout := time.Second * time.Duration(1800 + rand.Intn(600)) // a little random factor in here to spread out worker deaths
	task := worker.Task{
		CodeName: route.CodeName,
		Payload:  string(jsonPayload),
		Timeout:  &timeout, // let's have these die quickly while testing
	}
	tasks := make([]worker.Task, 1)
	tasks[0] = task
	taskIds, err := workerapi.TaskQueue(tasks...)
	golog.Infoln("Tasks queued.", taskIds)
	if err != nil {
		golog.Infoln("Couldn't queue up worker!", err)
		return err
	}
	return err
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

//	s, err := ioutil.ReadAll(req.Body)
//	fmt.Println("req.body:", err, string(s))

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	golog.Infoln("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		if !common.ReadJSON(w, req, &route) {
			return
		}
		golog.Infoln("body read into route:", route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		err := putRoute(&route)
		if err != nil {
			golog.Infoln("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		golog.Infoln("registered route:", route)
		fmt.Fprintln(w, "Host registered successfully.")

	} else {
		r2 := Route2{}
		if !common.ReadJSON(w, req, &r2) {
			return
		}
		// todo: do we need to close body?
		golog.Infoln("DECODED:", r2)
		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln("This host is not registered!", err))
			return
			//			route = &Route{}
		}
		golog.Infoln("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		golog.Infoln("ROUTE new:", route)
		err = putRoute(route)
		if err != nil {
			golog.Infoln("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (*Route, error) {
	rx, err := icache.Get(host)
	if err != nil {
		return nil, err
	}
	rx2 := []byte(rx.(string))
	route := Route{}
	err = json.Unmarshal(rx2, &route)
	if err != nil {
		return nil, err
	}
	return &route, err
}

func putRoute(route *Route) (error) {
	item := cache.Item{}
	v, err := json.Marshal(route)
	if err != nil {
		return err
	}
	item.Value = string(v)
	err = icache.Put(route.Host, &item)
	return err
}

func Ping(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintln(w, "pong")
}
<|endoftext|>


<commit_before>package hypercat

import (
	"encoding/json"
	"errors"
)

const (
	// HyperCatVersion is the version of HyperCat this library currently supports
	HyperCatVersion = "2.0"

	// MediaType is the default mime type of HyperCat resources
	HyperCatMediaType = "application/vnd.hypercat.catalogue+json"

	// DescriptionRel is the URI for the hasDescription relationship
	DescriptionRel = "urn:X-hypercat:rels:hasDescription:en"

	// ContentTypeRel is the URI for the isContentType relationship
	ContentTypeRel = "urn:X-hypercat:rels:isContentType"

	// HomepageRel is the URI for hasHomepage relationship
	HomepageRel = "urn:X-hypercat:rels:hasHomepage"

	// ContainsContentTypeRel is the URI for the containsContentType relationship
	ContainsContentTypeRel = "urn:X-hypercat:rels:containsContentType"

	// SupportsSearchRel is the URI for the supportsSearch relationship
	SupportsSearchRel = "urn:X-hypercat:rels:supportsSearch"
)

/*
 * HyperCat is the representation of the HyperCat catalogue object, which is
 * the parent element of each catalogue instance.
 */
type HyperCat struct {
	Items       Items    `json:"items"`
	Metadata    Metadata `json:"item-metadata"`
	Description string   `json:"-"` // HyperCat spec is fuzzy about whether there can be more than one description. We assume not.
}

/*
 * NewHyperCat is a constructor function that creates and returns a HyperCat
 * instance.
 */
func NewHyperCat(description string) *HyperCat {
	return &HyperCat{
		Description: description,
		Metadata:    Metadata{},
	}
}

/*
 * Parse is a convenience function that parses a HyperCat catalogue string, and
 * builds an in memory HyperCat instance.
 */
func Parse(str string) (*HyperCat, error) {
	cat := HyperCat{}
	err := json.Unmarshal([]byte(str), &cat)

	if err != nil {
		return nil, err
	}

	return &cat, nil
}

/*
 * AddItem is a convenience function for adding an Item to a catalogue.
 */
func (h *HyperCat) AddItem(item *Item) {
	h.Items = append(h.Items, *item)
}

/*
 * MarshalJSON returns the JSON encoding of a HyperCat. This function is the
 * implementation of the Marshaler interface.
 */
func (h *HyperCat) MarshalJSON() ([]byte, error) {
	metadata := h.Metadata

	if h.Description != "" {
		metadata = append(metadata, Rel{Rel: DescriptionRel, Val: h.Description})
	}

	return json.Marshal(struct {
		Items    []Item   `json:"items"`
		Metadata Metadata `json:"item-metadata"`
	}{
		Items:    h.Items,
		Metadata: metadata,
	})
}

/*
 * UnmarshalJSON is the required function for structs that implement the
 * Unmarshaler interface.
 */
func (h *HyperCat) UnmarshalJSON(b []byte) error {
	type tempCat struct {
		Items    Items    `json:"items"`
		Metadata Metadata `json:"item-metadata"`
	}

	t := tempCat{}

	err := json.Unmarshal(b, &t)

	if err != nil {
		return err
	}

	for _, rel := range t.Metadata {
		if rel.Rel == DescriptionRel {
			h.Description = rel.Val
		} else {
			h.Metadata = append(h.Metadata, rel)
		}
	}

	if h.Description == "" {
		err := errors.New(`"` + DescriptionRel + `" is a mandatory metadata element`)
		return err
	}

	return nil
}
<commit_msg>Some comment clean up.<commit_after>package hypercat

import (
	"encoding/json"
	"errors"
)

const (
	// HyperCatVersion is the version of HyperCat this library currently supports
	HyperCatVersion = "2.0"

	// HyperCatMediaType is the default mime type of HyperCat resources
	HyperCatMediaType = "application/vnd.hypercat.catalogue+json"

	// DescriptionRel is the URI for the hasDescription relationship
	DescriptionRel = "urn:X-hypercat:rels:hasDescription:en"

	// ContentTypeRel is the URI for the isContentType relationship
	ContentTypeRel = "urn:X-hypercat:rels:isContentType"

	// HomepageRel is the URI for hasHomepage relationship
	HomepageRel = "urn:X-hypercat:rels:hasHomepage"

	// ContainsContentTypeRel is the URI for the containsContentType relationship
	ContainsContentTypeRel = "urn:X-hypercat:rels:containsContentType"

	// SupportsSearchRel is the URI for the supportsSearch relationship
	SupportsSearchRel = "urn:X-hypercat:rels:supportsSearch"
)

/*
 * HyperCat is the representation of the HyperCat catalogue object, which is
 * the parent element of each catalogue instance.
 */
type HyperCat struct {
	Items       Items    `json:"items"`
	Metadata    Metadata `json:"item-metadata"`
	Description string   `json:"-"` // HyperCat spec is fuzzy about whether there can be more than one description. We assume not.
}

/*
 * NewHyperCat is a constructor function that creates and returns a HyperCat
 * instance.
 */
func NewHyperCat(description string) *HyperCat {
	return &HyperCat{
		Description: description,
		Metadata:    Metadata{},
	}
}

/*
 * Parse is a function that parses a HyperCat catalogue string, and builds an
 * in memory HyperCat instance.
 */
func Parse(str string) (*HyperCat, error) {
	cat := HyperCat{}
	err := json.Unmarshal([]byte(str), &cat)

	if err != nil {
		return nil, err
	}

	return &cat, nil
}

/*
 * AddItem is a convenience function for adding an Item to a catalogue.
 */
func (h *HyperCat) AddItem(item *Item) {
	h.Items = append(h.Items, *item)
}

/*
 * MarshalJSON returns the JSON encoding of a HyperCat. This function is the
 * implementation of the Marshaler interface.
 */
func (h *HyperCat) MarshalJSON() ([]byte, error) {
	metadata := h.Metadata

	if h.Description != "" {
		metadata = append(metadata, Rel{Rel: DescriptionRel, Val: h.Description})
	}

	return json.Marshal(struct {
		Items    []Item   `json:"items"`
		Metadata Metadata `json:"item-metadata"`
	}{
		Items:    h.Items,
		Metadata: metadata,
	})
}

/*
 * UnmarshalJSON is the required function for structs that implement the
 * Unmarshaler interface.
 */
func (h *HyperCat) UnmarshalJSON(b []byte) error {
	type tempCat struct {
		Items    Items    `json:"items"`
		Metadata Metadata `json:"item-metadata"`
	}

	t := tempCat{}

	err := json.Unmarshal(b, &t)

	if err != nil {
		return err
	}

	for _, rel := range t.Metadata {
		if rel.Rel == DescriptionRel {
			h.Description = rel.Val
		} else {
			h.Metadata = append(h.Metadata, rel)
		}
	}

	if h.Description == "" {
		err := errors.New(`"` + DescriptionRel + `" is a mandatory metadata element`)
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"sync/atomic"
	"time"
	"github.com/buger/gor/proto"
)

const initialDynamicWorkers = 10

type response struct {
	payload       []byte
	uuid          []byte
	roundTripTime int64
}

// HTTPOutputConfig struct for holding http output configuration
type HTTPOutputConfig struct {
	redirectLimit int

	stats   bool
	workers int

	elasticSearch string

	Timeout      time.Duration
	OriginalHost bool

	Debug bool

	TrackResponses bool
}

// HTTPOutput plugin manage pool of workers which send request to replayed server
// By default workers pool is dynamic and starts with 10 workers
// You can specify fixed number of workers using `--output-http-workers`
type HTTPOutput struct {
	// Keep this as first element of struct because it guarantees 64bit
	// alignment. atomic.* functions crash on 32bit machines if operand is not
	// aligned at 64bit. See https://github.com/golang/go/issues/599
	activeWorkers int64

	address string
	limit   int
	queue   chan []byte

	responses chan response

	needWorker chan int

	config *HTTPOutputConfig

	queueStats *GorStat

	elasticSearch *ESPlugin
}

// NewHTTPOutput constructor for HTTPOutput
// Initialize workers
func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
	o := new(HTTPOutput)

	o.address = address
	o.config = config

	if o.config.stats {
		o.queueStats = NewGorStat("output_http")
	}

	o.queue = make(chan []byte, 100)
	o.responses = make(chan response, 100)
	o.needWorker = make(chan int, 1)

	// Initial workers count
	if o.config.workers == 0 {
		o.needWorker <- initialDynamicWorkers
	} else {
		o.needWorker <- o.config.workers
	}

	if o.config.elasticSearch != "" {
		o.elasticSearch = new(ESPlugin)
		o.elasticSearch.Init(o.config.elasticSearch)
	}

	if len(Settings.middleware) > 0 {
		o.config.TrackResponses = true
	}

	go o.workerMaster()

	return o
}

func (o *HTTPOutput) workerMaster() {
	for {
		newWorkers := <-o.needWorker
		for i := 0; i < newWorkers; i++ {
			go o.startWorker()
		}

		// Disable dynamic scaling if workers poll fixed size
		if o.config.workers != 0 {
			return
		}
	}
}

func (o *HTTPOutput) startWorker() {
	client := NewHTTPClient(o.address, &HTTPClientConfig{
		FollowRedirects: o.config.redirectLimit,
		Debug:           o.config.Debug,
		OriginalHost:    o.config.OriginalHost,
		Timeout:         o.config.Timeout,
	})

	deathCount := 0

	atomic.AddInt64(&o.activeWorkers, 1)

	for {
		select {
		case data := <-o.queue:
			o.sendRequest(client, data)
			deathCount = 0
		case <-time.After(time.Millisecond * 100):
			// When dynamic scaling enabled workers die after 2s of inactivity
			if o.config.workers == 0 {
				deathCount++
			} else {
				continue
			}

			if deathCount > 20 {
				workersCount := atomic.LoadInt64(&o.activeWorkers)

				// At least 1 startWorker should be alive
				if workersCount != 1 {
					atomic.AddInt64(&o.activeWorkers, -1)
					return
				}
			}
		}
	}
}

func (o *HTTPOutput) Write(data []byte) (n int, err error) {
	if !isRequestPayload(data) {
		return len(data), nil
	}

	buf := make([]byte, len(data))
	copy(buf, data)

	o.queue <- buf

	if o.config.stats {
		o.queueStats.Write(len(o.queue))
	}

	if o.config.workers == 0 {
		workersCount := atomic.LoadInt64(&o.activeWorkers)

		if len(o.queue) > int(workersCount) {
			o.needWorker <- len(o.queue)
		}
	}

	return len(data), nil
}

func (o *HTTPOutput) Read(data []byte) (int, error) {
	resp := <-o.responses

	Debug("[OUTPUT-HTTP] Received response:", string(resp.payload))

	header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime)
	copy(data[0:len(header)], header)
	copy(data[len(header):], resp.payload)

	return len(resp.payload) + len(header), nil
}

func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) {
	meta := payloadMeta(request)
	uuid := meta[1]

	body := payloadBody(request)
	if !proto.IsHTTPPayload(body) {
		return
	}

	start := time.Now()
	resp, err := client.Send(body)
	stop := time.Now()

	if err != nil {
		log.Println("Request error:", err)
	}

	if o.config.TrackResponses {
		o.responses <- response{resp, uuid, stop.UnixNano() - start.UnixNano()}
	}

	if o.elasticSearch != nil {
		o.elasticSearch.ResponseAnalyze(request, resp, start, stop)
	}
}

func (o *HTTPOutput) String() string {
	return "HTTP output: " + o.address
}
<commit_msg>small formatting change<commit_after>package main

import (
	"github.com/buger/gor/proto"
	"io"
	"log"
	"sync/atomic"
	"time"
)

const initialDynamicWorkers = 10

type response struct {
	payload       []byte
	uuid          []byte
	roundTripTime int64
}

// HTTPOutputConfig struct for holding http output configuration
type HTTPOutputConfig struct {
	redirectLimit int

	stats   bool
	workers int

	elasticSearch string

	Timeout      time.Duration
	OriginalHost bool

	Debug bool

	TrackResponses bool
}

// HTTPOutput plugin manage pool of workers which send request to replayed server
// By default workers pool is dynamic and starts with 10 workers
// You can specify fixed number of workers using `--output-http-workers`
type HTTPOutput struct {
	// Keep this as first element of struct because it guarantees 64bit
	// alignment. atomic.* functions crash on 32bit machines if operand is not
	// aligned at 64bit. See https://github.com/golang/go/issues/599
	activeWorkers int64

	address string
	limit   int
	queue   chan []byte

	responses chan response

	needWorker chan int

	config *HTTPOutputConfig

	queueStats *GorStat

	elasticSearch *ESPlugin
}

// NewHTTPOutput constructor for HTTPOutput
// Initialize workers
func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
	o := new(HTTPOutput)

	o.address = address
	o.config = config

	if o.config.stats {
		o.queueStats = NewGorStat("output_http")
	}

	o.queue = make(chan []byte, 100)
	o.responses = make(chan response, 100)
	o.needWorker = make(chan int, 1)

	// Initial workers count
	if o.config.workers == 0 {
		o.needWorker <- initialDynamicWorkers
	} else {
		o.needWorker <- o.config.workers
	}

	if o.config.elasticSearch != "" {
		o.elasticSearch = new(ESPlugin)
		o.elasticSearch.Init(o.config.elasticSearch)
	}

	if len(Settings.middleware) > 0 {
		o.config.TrackResponses = true
	}

	go o.workerMaster()

	return o
}

func (o *HTTPOutput) workerMaster() {
	for {
		newWorkers := <-o.needWorker
		for i := 0; i < newWorkers; i++ {
			go o.startWorker()
		}

		// Disable dynamic scaling if workers poll fixed size
		if o.config.workers != 0 {
			return
		}
	}
}

func (o *HTTPOutput) startWorker() {
	client := NewHTTPClient(o.address, &HTTPClientConfig{
		FollowRedirects: o.config.redirectLimit,
		Debug:           o.config.Debug,
		OriginalHost:    o.config.OriginalHost,
		Timeout:         o.config.Timeout,
	})

	deathCount := 0

	atomic.AddInt64(&o.activeWorkers, 1)

	for {
		select {
		case data := <-o.queue:
			o.sendRequest(client, data)
			deathCount = 0
		case <-time.After(time.Millisecond * 100):
			// When dynamic scaling enabled workers die after 2s of inactivity
			if o.config.workers == 0 {
				deathCount++
			} else {
				continue
			}

			if deathCount > 20 {
				workersCount := atomic.LoadInt64(&o.activeWorkers)

				// At least 1 startWorker should be alive
				if workersCount != 1 {
					atomic.AddInt64(&o.activeWorkers, -1)
					return
				}
			}
		}
	}
}

func (o *HTTPOutput) Write(data []byte) (n int, err error) {
	if !isRequestPayload(data) {
		return len(data), nil
	}

	buf := make([]byte, len(data))
	copy(buf, data)

	o.queue <- buf

	if o.config.stats {
		o.queueStats.Write(len(o.queue))
	}

	if o.config.workers == 0 {
		workersCount := atomic.LoadInt64(&o.activeWorkers)

		if len(o.queue) > int(workersCount) {
			o.needWorker <- len(o.queue)
		}
	}

	return len(data), nil
}

func (o *HTTPOutput) Read(data []byte) (int, error) {
	resp := <-o.responses

	Debug("[OUTPUT-HTTP] Received response:", string(resp.payload))

	header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime)
	copy(data[0:len(header)], header)
	copy(data[len(header):], resp.payload)

	return len(resp.payload) + len(header), nil
}

func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) {
	meta := payloadMeta(request)
	uuid := meta[1]

	body := payloadBody(request)
	if !proto.IsHTTPPayload(body) {
		return
	}

	start := time.Now()
	resp, err := client.Send(body)
	stop := time.Now()

	if err != nil {
		log.Println("Request error:", err)
	}

	if o.config.TrackResponses {
		o.responses <- response{resp, uuid, stop.UnixNano() - start.UnixNano()}
	}

	if o.elasticSearch != nil {
		o.elasticSearch.ResponseAnalyze(request, resp, start, stop)
	}
}

func (o *HTTPOutput) String() string {
	return "HTTP output: " + o.address
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"regexp"
	"strings"
	"fmt"
)

func Identify(filename string) {
	content, err := ioutil.ReadFile(filename)
	ChkErr(err)
	text := string(content)

	// example: re, err := regexp.Compile(`^\[[ ]?(\d{1,2}) ([[:alpha:]-]{3})\] ([[:alpha:]]+) .*\((.*)\)`)

	// put all flags/restricts, or effects, on one line
	re, err := regexp.Compile(`([[:upper:]]{2})\n([[:upper:]]{2})`)
	ChkErr(err)
	text = re.ReplaceAllString(text, "$1 $2")

	// put enchant info on one line
	re, err = regexp.Compile(`\n(Duration)`)
	ChkErr(err)
	text = re.ReplaceAllString(text, " $1")

	// put all resists on same line:
	re, err = regexp.Compile(`\n(    [[:alpha:] ]{6}:[[:blank:]]{3,4}[[:digit:]]{1,2}%)`)
	ChkErr(err)
	text = re.ReplaceAllString(text, "$1")

	// put first wand/staff/scroll/potion spell on one line:
	re, err = regexp.Compile(`(spell[s]? of:)\n`)
	ChkErr(err)
	text = re.ReplaceAllString(text, "$1 ")

	// put remaining scroll/potion spells on same line:
	re, err = regexp.Compile(`\n([[:lower:]])`)
	ChkErr(err)
	text = re.ReplaceAllString(text, ", $1")

	fmt.Println(text) // debugging

	items := strings.Split(text, "\n\n")

	for _, item := range items {
		lines := strings.Split(item, "\n")

		for _, line := range lines {
			// use regex to capture useful info
			_ = line
		}
	}
}
<commit_msg>Additional regex, starting to pull items apart<commit_after>package main

import (
	"io/ioutil"
	"regexp"
	"strings"
	"fmt"
)

func Identify(filename string) {
	content, err := ioutil.ReadFile(filename)
	ChkErr(err)
	text := string(content)

	// put all flags/restricts, or effects, on one line
	re, err := regexp.Compile(`([[:upper:]]{2})\n([[:upper:]]{2})`)
	ChkErr(err)
	text = re.ReplaceAllString(text, "$1 $2")

	// put enchant info on one line
	re, err = regexp.Compile(`\n(Duration)`)
	ChkErr(err)
	text = re.ReplaceAllString(text, " $1")

	// put all resists on same line:
	re, err = regexp.Compile(
		`\n(    [[:alpha:] ]{6}:[[:blank:]]{3,4}[[:digit:]]{1,2}% )`)
	ChkErr(err)
	text = re.ReplaceAllString(text, "$1")

	// put first wand/staff/scroll/potion spell on one line:
	re, err = regexp.Compile(`(spell[s]? of:)\n`)
	ChkErr(err)
	text = re.ReplaceAllString(text, "$1 ")

	// put remaining scroll/potion spells on same line:
	re, err = regexp.Compile(`\n([[:lower:]])`)
	ChkErr(err)
	text = re.ReplaceAllString(text, ", $1")

	items := strings.Split(text, "\n\n")

	// initialize item variables
	full_stats, item_name, keywords, item_type, slot := "", "", "", "", ""

	// initialize regex checks
	var m [][]string
	chkName, err := regexp.Compile(
		`Name '([[:print:]]+)'`)
	ChkErr(err)
	chkKey, err := regexp.Compile(
		`Keyword '([[:print:]]+)', Item type: ([[:word:]]+)`)
	ChkErr(err)
	chkWorn, err := regexp.Compile(
		`Item can be worn on:  ([[:print:]]+) `)
	ChkErr(err)

	for _, item := range items {
		full_stats = item
		lines := strings.Split(item, "\n")

		for _, line := range lines {
			switch {
			case chkName.MatchString(line):
				m = chkName.FindAllStringSubmatch(line, -1)
				item_name = m[0][1]
			case chkKey.MatchString(line):
				m = chkKey.FindAllStringSubmatch(line, -1)
				keywords = m[0][1]
				item_type = m[0][2]
			case chkWorn.MatchString(line):
				m = chkWorn.FindAllStringSubmatch(line, -1)
				slot = m[0][1]
			}
		}
		// back to item
		fmt.Printf("Name: %s, Keywords: %s, Type: %s, Slot: %s\n",
			item_name, keywords, item_type, slot)
		_ = full_stats
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 Aaron Donovan <amdonov@gmail.com>
//
// 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 idp

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"encoding/xml"
	"io"
	"net/http"

	"github.com/amdonov/lite-idp/model"
)

func (i *IDP) sendPostResponse(authRequest *model.AuthnRequest, user *model.User,
	w io.Writer, r *http.Request) error {
	response := i.makeAuthnResponse(authRequest, user)
	// Don't need to change the response. Go ahead and sign it
	signature, err := i.signer.CreateSignature(response.Assertion)
	if err != nil {
		return err
	}
	response.Assertion.Signature = signature
	var xmlbuff bytes.Buffer
	memWriter := bufio.NewWriter(&xmlbuff)
	memWriter.Write([]byte(xml.Header))
	encoder := xml.NewEncoder(memWriter)
	encoder.Encode(response)
	memWriter.Flush()

	samlMessage := base64.StdEncoding.EncodeToString(xmlbuff.Bytes())

	data := struct {
		RelayState                  string
		SAMLResponse                string
		AssertionConsumerServiceURL string
	}{
		authRequest.RelayState,
		samlMessage,
		authRequest.AssertionConsumerServiceURL,
	}
	return i.postTemplate.Execute(w, data)
}

//Assume HTML 5, where  is not required
const postTemplate = `<!DOCTYPE html>

Note: Since your browser does not support JavaScript,
you must press the Continue button once to proceed.

<form action="{{ .AssertionConsumerServiceURL }}" method="post" id="samlpost">

<input type="hidden" name="RelayState"
value="{{ .RelayState }}"/>
<input type="hidden" name="SAMLResponse"
value="{{ .SAMLResponse }}"/>

<input type="submit" value="Continue"/>

</form>

`
<commit_msg>Remove xmlns attribute which is not part of html5 spec<commit_after>// Copyright © 2017 Aaron Donovan <amdonov@gmail.com>
//
// 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 idp

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"encoding/xml"
	"io"
	"net/http"

	"github.com/amdonov/lite-idp/model"
)

func (i *IDP) sendPostResponse(authRequest *model.AuthnRequest, user *model.User,
	w io.Writer, r *http.Request) error {
	response := i.makeAuthnResponse(authRequest, user)
	// Don't need to change the response. Go ahead and sign it
	signature, err := i.signer.CreateSignature(response.Assertion)
	if err != nil {
		return err
	}
	response.Assertion.Signature = signature
	var xmlbuff bytes.Buffer
	memWriter := bufio.NewWriter(&xmlbuff)
	memWriter.Write([]byte(xml.Header))
	encoder := xml.NewEncoder(memWriter)
	encoder.Encode(response)
	memWriter.Flush()

	samlMessage := base64.StdEncoding.EncodeToString(xmlbuff.Bytes())

	data := struct {
		RelayState                  string
		SAMLResponse                string
		AssertionConsumerServiceURL string
	}{
		authRequest.RelayState,
		samlMessage,
		authRequest.AssertionConsumerServiceURL,
	}
	return i.postTemplate.Execute(w, data)
}

//Assume HTML 5, where  is not required
const postTemplate = `<!DOCTYPE html>

Note: Since your browser does not support JavaScript,
you must press the Continue button once to proceed.

<form action="{{ .AssertionConsumerServiceURL }}" method="post" id="samlpost">

<input type="hidden" name="RelayState"
value="{{ .RelayState }}"/>
<input type="hidden" name="SAMLResponse"
value="{{ .SAMLResponse }}"/>

<input type="submit" value="Continue"/>

</form>

`
<|endoftext|>