<commit_before>package hotp

import (
	"bytes"
	"code.google.com/p/rsc/qr"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"math/big"
	"net/url"
	"strconv"
)

// RFC 4226 specifies the counter as being 8 bytes.
const ctrSize = 8

var (
	ErrInvalidHOTPURL = errors.New("hotp: invalid HOTP url")
)

// Type HOTP represents a new source for generating one-time passwords.
type HOTP struct {
	Key     []byte
	counter *[ctrSize]byte
	Digits  int
}

// Counter returns the HOTP's 8-byte counter as an unsigned 64-bit
// integer.
func (otp HOTP) Counter() uint64 {
	buf := bytes.NewBuffer(otp.counter[:])
	var counter uint64
	err := binary.Read(buf, binary.BigEndian, &counter)
	if err != nil {
		panic("counter should never be invalid")
	}
	return counter
}

// Increment will increment an HOTP source's counter. This is useful
// for providers like the Google Authenticator app, which immediately
// increments the counter and uses the 0 counter value as an integrity
// check.
func (otp HOTP) Increment() {
	for i := ctrSize - 1; i >= 0; i-- {
		if otp.counter[i]++; otp.counter[i] != 0 {
			return
		}
	}
}

// OTP generates a new one-time password.
func (otp HOTP) OTP() string {
	h := hmac.New(sha1.New, otp.Key)
	h.Write(otp.counter[:])
	otp.Increment()
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64())
}

func (otp *HOTP) setCounter(counter uint64) bool {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, counter)
	if err != nil {
		return false
	}

	ctr := zeroPad(buf.Bytes())
	if ctr == nil {
		return false
	}
	var ctr8 [ctrSize]byte
	copy(ctr8[:], ctr)
	copy(otp.counter[:], ctr8[:])
	return true
}

// NewHOTP intialises a new HOTP instance with the key and counter
// values. No check is done on the digits, but typical values are 6
// and 8.
func NewHOTP(key []byte, counter uint64, digits int) *HOTP {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, counter)
	if err != nil {
		return nil
	}

	ctr := zeroPad(buf.Bytes())
	if ctr == nil {
		return nil
	}
	var ctr8 [ctrSize]byte
	copy(ctr8[:], ctr)
	otp := &HOTP{
		Key:    key,
		Digits: digits,
	}
	otp.counter = new([ctrSize]byte)
	copy(otp.counter[:], ctr8[:])
	return otp
}

// URL returns a suitable URL, such as for the Google Authenticator
// app. The label is used by these apps to identify the service to
// which this OTP belongs. The digits value is ignored by the Google
// authenticator app, and is therefore elided in the resulting URL.
func (otp *HOTP) URL(label string) string {
	secret := base32.StdEncoding.EncodeToString(otp.Key)
	u := url.URL{}
	v := url.Values{}
	u.Scheme = "otpauth"
	u.Host = "hotp"
	u.Path = label
	v.Add("secret", secret)
	v.Add("counter", fmt.Sprintf("%d", otp.Counter()))
	u.RawQuery = v.Encode()
	return u.String()
}

// QR generates a byte slice containing the a QR code encoded as a
// PNG with level Q error correction.
func (otp *HOTP) QR(label string) ([]byte, error) {
	u := otp.URL(label)
	code, err := qr.Encode(u, qr.Q)
	if err != nil {
		return nil, err
	}
	return code.PNG(), nil
}

func zeroPad(in []byte) []byte {
	inLen := len(in)
	if inLen > ctrSize {
		return in[:8]
	}
	start := ctrSize - inLen
	out := make([]byte, ctrSize)
	copy(out[start:], in)
	return out
}

func truncate(in []byte) int64 {
	offset := int(in[len(in)-1] & 0xF)
	p := in[offset : offset+4]
	var binCode int32
	binCode = int32((p[0] & 0x7f)) << 24
	binCode += int32((p[1] & 0xff)) << 16
	binCode += int32((p[2] & 0xff)) << 8
	binCode += int32((p[3] & 0xff))
	return int64(binCode) & 0x7FFFFFFF
}

// FromURL parses a new HOTP from a URL string. It returns the OTP,
// the label associated with the OTP, and any errors that occurred.
func FromURL(urlString string) (*HOTP, string, error) {
	u, err := url.Parse(urlString)
	if err != nil {
		return nil, "", err
	}

	if u.Scheme != "otpauth" {
		return nil, "", ErrInvalidHOTPURL
	} else if u.Host != "hotp" {
		return nil, "", ErrInvalidHOTPURL
	}

	v := u.Query()
	if len(v) == 0 {
		v, err = url.ParseQuery(u.Path[1:])
		if err != nil {
			return nil, "", err
		}

	}
	if v.Get("secret") == "" {
		return nil, "", ErrInvalidHOTPURL
	} else if algo := v.Get("algorithm"); algo != "" && algo != "SHA1" {
		return nil, "", ErrInvalidHOTPURL
	}

	var identity string
	if len(u.Path) > 1 {
		identity = u.Path[1:]
	}

	var counter uint64
	if ctr := v.Get("counter"); ctr != "" {
		counter, err = strconv.ParseUint(ctr, 10, 64)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	secret, err := base32.StdEncoding.DecodeString(v.Get("secret"))
	if err != nil {
		return nil, "", ErrInvalidHOTPURL
	}

	var digits int64 = 6
	if v.Get("digits") != "" {
		digits, err = strconv.ParseInt(v.Get("digits"), 10, 8)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	otp := NewHOTP(secret, counter, int(digits))
	return otp, identity, nil
}

// GenerateHOTP will generate a randomised HOTP source; if the
// randCounter parameter is true, the counter will be randomised.
func GenerateHOTP(digits int, randCounter bool) (*HOTP, error) {
	key := make([]byte, sha1.Size)
	_, err := io.ReadFull(rand.Reader, key)
	if err != nil {
		return nil, err
	}

	var counter uint64
	if randCounter {
		ctr, err := rand.Int(rand.Reader, big.NewInt(int64(math.MaxInt64)))
		if err != nil {
			return nil, err
		}
		counter = ctr.Uint64()
	}

	return NewHOTP(key, counter, digits), nil
}

// YubiKey reads an OATH-HOTP string as returned by a YubiKey, and
// returns three values. The first value contains the actual OTP, the
// second value contains the YubiKey's token identifier, and the final
// value indicates whether the input string was a valid YubiKey
// OTP. This does not check whether the code is correct or not, it
// only ensures that it is well-formed output from a token and
// splits the output into the code and the public identity.
func (otp *HOTP) YubiKey(in string) (string, string, bool) {
	if len(in) < otp.Digits {
		return "", "", false
	}

	otpStart := len(in) - otp.Digits
	code := in[otpStart:]
	pubid := in[:otpStart]
	return code, pubid, true
}

// IntegrityCheck returns two values, the base OTP and the current
// counter. This is used, for example, with the Google Authenticator
// app's "Check key value" function and can be used to verify that
// the application and the provider are in sync.
func (otp *HOTP) IntegrityCheck() (string, uint64) {
	h := hmac.New(sha1.New, otp.Key)
	counter := make([]byte, 8)
	h.Write(counter)
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64()), otp.Counter()
}

// Scan takes a code input (i.e. from the user), and scans ahead
// within a certain window of counter values. This can be used in the
// case where the server's counter and the user's counter have fallen
// out of sync.
func (otp *HOTP) Scan(code string, window int) bool {
	var valid bool
	codeBytes := []byte(code)
	counter := otp.Counter()

	for i := 0; i < window; i++ {
		genCode := []byte(otp.OTP())
		if subtle.ConstantTimeCompare(codeBytes, genCode) == 1 {
			valid = true
			break
		}
	}
	if !valid {
		otp.setCounter(counter)
	}
	return valid
}

// Check takes an input code and verifies it against the OTP. If
// successful, the counter is incremented.
func (otp *HOTP) Check(code string) bool {
	codeBytes := []byte(code)
	genCode := []byte(otp.OTP())
	if subtle.ConstantTimeCompare(codeBytes, genCode) != 1 {
		otp.setCounter(otp.Counter() - 1)
		return false
	} else {
		return true
	}
}
<commit_msg>Clean up variable declaration.<commit_after>package hotp

import (
	"bytes"
	"code.google.com/p/rsc/qr"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"math/big"
	"net/url"
	"strconv"
)

// RFC 4226 specifies the counter as being 8 bytes.
const ctrSize = 8

var ErrInvalidHOTPURL = errors.New("hotp: invalid HOTP url")

// Type HOTP represents a new source for generating one-time passwords.
type HOTP struct {
	Key     []byte
	counter *[ctrSize]byte
	Digits  int
}

// Counter returns the HOTP's 8-byte counter as an unsigned 64-bit
// integer.
func (otp HOTP) Counter() uint64 {
	buf := bytes.NewBuffer(otp.counter[:])
	var counter uint64
	err := binary.Read(buf, binary.BigEndian, &counter)
	if err != nil {
		panic("counter should never be invalid")
	}
	return counter
}

// Increment will increment an HOTP source's counter. This is useful
// for providers like the Google Authenticator app, which immediately
// increments the counter and uses the 0 counter value as an integrity
// check.
func (otp HOTP) Increment() {
	for i := ctrSize - 1; i >= 0; i-- {
		if otp.counter[i]++; otp.counter[i] != 0 {
			return
		}
	}
}

// OTP generates a new one-time password.
func (otp HOTP) OTP() string {
	h := hmac.New(sha1.New, otp.Key)
	h.Write(otp.counter[:])
	otp.Increment()
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64())
}

func (otp *HOTP) setCounter(counter uint64) bool {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, counter)
	if err != nil {
		return false
	}

	ctr := zeroPad(buf.Bytes())
	if ctr == nil {
		return false
	}
	var ctr8 [ctrSize]byte
	copy(ctr8[:], ctr)
	copy(otp.counter[:], ctr8[:])
	return true
}

// NewHOTP intialises a new HOTP instance with the key and counter
// values. No check is done on the digits, but typical values are 6
// and 8.
func NewHOTP(key []byte, counter uint64, digits int) *HOTP {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, counter)
	if err != nil {
		return nil
	}

	ctr := zeroPad(buf.Bytes())
	if ctr == nil {
		return nil
	}
	var ctr8 [ctrSize]byte
	copy(ctr8[:], ctr)
	otp := &HOTP{
		Key:    key,
		Digits: digits,
	}
	otp.counter = new([ctrSize]byte)
	copy(otp.counter[:], ctr8[:])
	return otp
}

// URL returns a suitable URL, such as for the Google Authenticator
// app. The label is used by these apps to identify the service to
// which this OTP belongs. The digits value is ignored by the Google
// authenticator app, and is therefore elided in the resulting URL.
func (otp *HOTP) URL(label string) string {
	secret := base32.StdEncoding.EncodeToString(otp.Key)
	u := url.URL{}
	v := url.Values{}
	u.Scheme = "otpauth"
	u.Host = "hotp"
	u.Path = label
	v.Add("secret", secret)
	v.Add("counter", fmt.Sprintf("%d", otp.Counter()))
	u.RawQuery = v.Encode()
	return u.String()
}

// QR generates a byte slice containing the a QR code encoded as a
// PNG with level Q error correction.
func (otp *HOTP) QR(label string) ([]byte, error) {
	u := otp.URL(label)
	code, err := qr.Encode(u, qr.Q)
	if err != nil {
		return nil, err
	}
	return code.PNG(), nil
}

func zeroPad(in []byte) []byte {
	inLen := len(in)
	if inLen > ctrSize {
		return in[:8]
	}
	start := ctrSize - inLen
	out := make([]byte, ctrSize)
	copy(out[start:], in)
	return out
}

func truncate(in []byte) int64 {
	offset := int(in[len(in)-1] & 0xF)
	p := in[offset : offset+4]
	var binCode int32
	binCode = int32((p[0] & 0x7f)) << 24
	binCode += int32((p[1] & 0xff)) << 16
	binCode += int32((p[2] & 0xff)) << 8
	binCode += int32((p[3] & 0xff))
	return int64(binCode) & 0x7FFFFFFF
}

// FromURL parses a new HOTP from a URL string. It returns the OTP,
// the label associated with the OTP, and any errors that occurred.
func FromURL(urlString string) (*HOTP, string, error) {
	u, err := url.Parse(urlString)
	if err != nil {
		return nil, "", err
	}

	if u.Scheme != "otpauth" {
		return nil, "", ErrInvalidHOTPURL
	} else if u.Host != "hotp" {
		return nil, "", ErrInvalidHOTPURL
	}

	v := u.Query()
	if len(v) == 0 {
		v, err = url.ParseQuery(u.Path[1:])
		if err != nil {
			return nil, "", err
		}

	}
	if v.Get("secret") == "" {
		return nil, "", ErrInvalidHOTPURL
	} else if algo := v.Get("algorithm"); algo != "" && algo != "SHA1" {
		return nil, "", ErrInvalidHOTPURL
	}

	var identity string
	if len(u.Path) > 1 {
		identity = u.Path[1:]
	}

	var counter uint64
	if ctr := v.Get("counter"); ctr != "" {
		counter, err = strconv.ParseUint(ctr, 10, 64)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	secret, err := base32.StdEncoding.DecodeString(v.Get("secret"))
	if err != nil {
		return nil, "", ErrInvalidHOTPURL
	}

	var digits int64 = 6
	if v.Get("digits") != "" {
		digits, err = strconv.ParseInt(v.Get("digits"), 10, 8)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	otp := NewHOTP(secret, counter, int(digits))
	return otp, identity, nil
}

// GenerateHOTP will generate a randomised HOTP source; if the
// randCounter parameter is true, the counter will be randomised.
func GenerateHOTP(digits int, randCounter bool) (*HOTP, error) {
	key := make([]byte, sha1.Size)
	_, err := io.ReadFull(rand.Reader, key)
	if err != nil {
		return nil, err
	}

	var counter uint64
	if randCounter {
		ctr, err := rand.Int(rand.Reader, big.NewInt(int64(math.MaxInt64)))
		if err != nil {
			return nil, err
		}
		counter = ctr.Uint64()
	}

	return NewHOTP(key, counter, digits), nil
}

// YubiKey reads an OATH-HOTP string as returned by a YubiKey, and
// returns three values. The first value contains the actual OTP, the
// second value contains the YubiKey's token identifier, and the final
// value indicates whether the input string was a valid YubiKey
// OTP. This does not check whether the code is correct or not, it
// only ensures that it is well-formed output from a token and
// splits the output into the code and the public identity.
func (otp *HOTP) YubiKey(in string) (string, string, bool) {
	if len(in) < otp.Digits {
		return "", "", false
	}

	otpStart := len(in) - otp.Digits
	code := in[otpStart:]
	pubid := in[:otpStart]
	return code, pubid, true
}

// IntegrityCheck returns two values, the base OTP and the current
// counter. This is used, for example, with the Google Authenticator
// app's "Check key value" function and can be used to verify that
// the application and the provider are in sync.
func (otp *HOTP) IntegrityCheck() (string, uint64) {
	h := hmac.New(sha1.New, otp.Key)
	counter := make([]byte, 8)
	h.Write(counter)
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64()), otp.Counter()
}

// Scan takes a code input (i.e. from the user), and scans ahead
// within a certain window of counter values. This can be used in the
// case where the server's counter and the user's counter have fallen
// out of sync.
func (otp *HOTP) Scan(code string, window int) bool {
	var valid bool
	codeBytes := []byte(code)
	counter := otp.Counter()

	for i := 0; i < window; i++ {
		genCode := []byte(otp.OTP())
		if subtle.ConstantTimeCompare(codeBytes, genCode) == 1 {
			valid = true
			break
		}
	}
	if !valid {
		otp.setCounter(counter)
	}
	return valid
}

// Check takes an input code and verifies it against the OTP. If
// successful, the counter is incremented.
func (otp *HOTP) Check(code string) bool {
	codeBytes := []byte(code)
	genCode := []byte(otp.OTP())
	if subtle.ConstantTimeCompare(codeBytes, genCode) != 1 {
		otp.setCounter(otp.Counter() - 1)
		return false
	} else {
		return true
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"bytes"
	"strings"
	"github.com/JalfResi/justext"
)

// Convert HTML
func ConvertHtml(input io.Reader, readability bool) (string, map[string]string) {
	meta := make(map[string]string)
	cleanXml, err := Tidy(input, false)
	if err != nil {
		log.Println("Tidy:", err)
	}
	if readability {
		cleanXml = HtmlReadability(bytes.NewReader(cleanXml))
	}
	return Html2Text(bytes.NewReader(cleanXml)), meta
}

// Extract the readable text in an HTML document
func HtmlReadability(input io.Reader) []byte {

	// Create a justext reader from another reader
	reader := justext.NewReader(input)

	// Configure the reader
	reader.Stoplist = readabilityStopList
	reader.LengthLow = *readabilityLengthLow
	reader.LengthHigh = *readabilityLengthHigh
	reader.StopwordsLow = *readabilityStopwordsLow
	reader.StopwordsHigh = *readabilityStopwordsHigh
	reader.MaxLinkDensity = *readabilityMaxLinkDensity
	reader.MaxHeadingDistance = *readabilityMaxHeadingDistance

	// Read from the reader to generate a paragraph set
	paragraphSet, err := reader.ReadAll()
	if err != nil {
		log.Println("Justext:", err)
		return nil
	}

	useClasses := strings.SplitN(*readabilityUseClasses, ",", 10)

	var output string = ""
	for _, paragraph := range paragraphSet {
		for _, class := range useClasses {
			if paragraph.CfClass == class {
				output += paragraph.Text+"\n"
			}
		}
	}

	return []byte(output)
}

func Html2Text(input io.Reader) string {
	return Xml2Text(input, []string{ "br", "p", "h1", "h2", "h3", "h4" }, []string{}, false)
}

var readabilityStopList map[string]bool = map[string]bool{ "and":true, "the":true, "a":true, "about":true, "above":true, "across":true, "after":true, "afterwards":true, "again":true, "against":true, "all":true, "almost":true, "alone":true,
	"along":true, "already":true, "also":true, "although":true, "always":true, "am":true, "among":true, "amongst":true, "amoungst":true, "amount":true, "an":true, "another":true, "any":true,
	"anyhow":true, "anyone":true, "anything":true, "anyway":true, "anywhere":true, "are":true, "around":true, "as":true, "at":true, "back":true, "be":true, "became":true, "because":true,
	"become":true, "becomes":true, "becoming":true, "been":true, "before":true, "beforehand":true, "behind":true, "being":true, "below":true, "beside":true, "besides":true, "between":true,
	"beyond":true, "both":true, "bottom":true, "but":true, "by":true, "can":true, "cannot":true, "cant":true, "co":true, "con":true, "could":true, "couldnt":true, "cry":true,
	"de":true, "describe":true, "detail":true, "do":true, "done":true, "down":true, "due":true, "during":true, "each":true, "eg":true, "eight":true, "either":true, "eleven":true, "else":true,
	"elsewhere":true, "empty":true, "enough":true, "etc":true, "even":true, "ever":true, "every":true, "everyone":true, "everything":true, "everywhere":true, "except":true, "few":true,
	"fifteen":true, "fify":true, "fill":true, "find":true, "fire":true, "first":true, "five":true, "for":true, "former":true, "formerly":true, "forty":true, "found":true, "four":true, "from":true,
	"front":true, "full":true, "further":true, "get":true, "give":true, "go":true, "had":true, "has":true, "hasnt":true, "have":true, "he":true, "hence":true, "her":true, "here":true, "hereafter":true,
	"hereby":true, "herein":true, "hereupon":true, "hers":true, "herself":true, "him":true, "himself":true, "his":true, "how":true, "however":true, "hundred":true, "ie":true, "if":true, "in":true,
	"inc":true, "indeed":true, "interest":true, "into":true, "is":true, "it":true, "its":true, "itself":true, "keep":true, "last":true, "latter":true, "latterly":true, "least":true, "less":true,
	"ltd":true, "made":true, "many":true, "may":true, "me":true, "meanwhile":true, "might":true, "mill":true, "mine":true, "more":true, "moreover":true, "most":true, "mostly":true, "move":true,
	"much":true, "must":true, "my":true, "myself":true, "name":true, "namely":true, "neither":true, "never":true, "nevertheless":true, "next":true, "nine":true, "no":true, "nobody":true,
	"none":true, "noone":true, "nor":true, "not":true, "nothing":true, "now":true, "nowhere":true, "of":true, "off":true, "often":true, "on":true, "once":true, "one":true, "only":true, "onto":true,
	"or":true, "other":true, "others":true, "otherwise":true, "our":true, "ours":true, "ourselves":true, "out":true, "over":true, "own":true, "part":true, "per":true, "perhaps":true,
	"please":true, "put":true, "rather":true, "re":true, "same":true, "see":true, "seem":true, "seemed":true, "seeming":true, "seems":true, "serious":true, "several":true, "she":true,
	"should":true, "show":true, "side":true, "since":true, "sincere":true, "six":true, "sixty":true, "so":true, "some":true, "somehow":true, "someone":true, "something":true, "sometime":true,
	"sometimes":true, "somewhere":true, "still":true, "such":true, "take":true, "ten":true, "than":true, "that":true, "their":true, "them":true, "themselves":true,
	"then":true, "thence":true, "there":true, "thereafter":true, "thereby":true, "therefore":true, "therein":true, "thereupon":true, "these":true, "they":true, "thickv":true, "thin":true,
	"third":true, "this":true, "those":true, "though":true, "three":true, "through":true, "throughout":true, "thru":true, "thus":true, "to":true, "together":true, "too":true, "top":true,
	"toward":true, "towards":true, "twelve":true, "twenty":true, "two":true, "un":true, "under":true, "until":true, "up":true, "upon":true, "us":true, "very":true, "via":true, "was":true, "we":true,
	"well":true, "were":true, "what":true, "whatever":true, "when":true, "whence":true, "whenever":true, "where":true, "whereafter":true, "whereas":true, "whereby":true, "wherein":true,
	"whereupon":true, "wherever":true, "whether":true, "which":true, "while":true, "whither":true, "who":true, "whoever":true, "whole":true, "whom":true, "whose":true, "why":true, "will":true,
	"with":true, "within":true, "without":true, "would":true, "yet":true, "you":true, "your":true, "youre":true, "yours":true, "yourself":true, "yourselves":true, "www":true, "com":true, "http":true }
<commit_msg>Add manual Golang tokenization to take over HTML cleaning when Tidy fails<commit_after>package main

import (
	"io"
	"log"
	"bytes"
	"strings"
	"github.com/JalfResi/justext"
	"code.google.com/p/go.net/html"
)

// Convert HTML
func ConvertHtml(input io.Reader, readability bool) (string, map[string]string) {
	meta := make(map[string]string)

	buf := new(bytes.Buffer)
	buf.ReadFrom(input)

	cleanXml, err := Tidy(bytes.NewReader(buf.Bytes()), false)
	if err != nil {
		log.Println("Tidy:", err)
		// Tidy failed, so we now manually tokenize instead
		clean := cleanHTML(buf, true)
		cleanXml = []byte(clean)
		log.Println("Cleaned HTML using Golang tokenizer")
	}

	if readability {
		cleanXml = HtmlReadability(bytes.NewReader(cleanXml))
	}
	return Html2Text(bytes.NewReader(cleanXml)), meta
}

// Tests for known friendly HTML parameters that tidy is unlikely to choke on
func acceptedHTML(str string) bool {
    tags := []string{"div", "p", "br", "span", "body", "head", "html", "ul", "ol", "li", "dl", "dt", "dd", "a", "form", "article", "table", "tr", "td", "tbody", "thead", "th", "tfoot", "col", "colgroup", "caption", "form", "input", "title", "h1", "h2", "h3", "h4", "h5", "h6", "meta", "strong", "cite", "em", "address", "abbr", "acronym", "blockquote", "q", "pre", "samp", "select", "fieldset", "legend", "button", "option", "textarea", "label"}
    for _, tag := range tags {
        if tag == str {
            return true
        }
    }
    return false
}

// Removes scripts, comments, styles and parameters from HTML.
// Also removes made up tags, e.g. <fb:like>
// Can keep head elements or not. Typically not much in there.
func cleanHTML(r io.Reader, all bool) string {
    d := html.NewTokenizer(r)

    output := ""
    if !all {
        output = "<html><head></head>"
    }
    mainsection := false
    junksection := false

    for {
        // token type
        tokenType := d.Next()
        if tokenType == html.ErrorToken {
            return output
        }
        token := d.Token()

        // type Token struct {
        //     Type     TokenType
        //     DataAtom atom.Atom
        //     Data     string
        //     Attr     []Attribute
        // }
        //
        // type Attribute struct {
        //     Namespace, Key, Val string
        // }

        switch tokenType {
            case html.StartTagToken: // <tag>
                if token.Data == "body" || (token.Data == "html" && all) {
                    mainsection = true
                }
                if !acceptedHTML(token.Data) {
                    junksection = true
                }

                if !junksection && mainsection {
                    output += "<" + token.Data + ">"
                }

            case html.TextToken: // text between start and end tag
                if !junksection && mainsection {
                    output += token.Data
                }

            case html.EndTagToken: // </tag>
                if !junksection && mainsection {
                    output += "</" + token.Data + ">"
                }
                if !acceptedHTML(token.Data) {
                    junksection = false
                }

            case html.SelfClosingTagToken: // <tag/>
                if !junksection && mainsection {
                    output += "<" + token.Data + " />"  // TODO: Can probably keep attributes from the meta tags
                }
        }
    }
}

// Extract the readable text in an HTML document
func HtmlReadability(input io.Reader) []byte {

	// Create a justext reader from another reader
	reader := justext.NewReader(input)

	// Configure the reader
	reader.Stoplist = readabilityStopList
	reader.LengthLow = *readabilityLengthLow
	reader.LengthHigh = *readabilityLengthHigh
	reader.StopwordsLow = *readabilityStopwordsLow
	reader.StopwordsHigh = *readabilityStopwordsHigh
	reader.MaxLinkDensity = *readabilityMaxLinkDensity
	reader.MaxHeadingDistance = *readabilityMaxHeadingDistance

	// Read from the reader to generate a paragraph set
	paragraphSet, err := reader.ReadAll()
	if err != nil {
		log.Println("Justext:", err)
		return nil
	}

	useClasses := strings.SplitN(*readabilityUseClasses, ",", 10)

	var output string = ""
	for _, paragraph := range paragraphSet {
		for _, class := range useClasses {
			if paragraph.CfClass == class {
				output += paragraph.Text+"\n"
			}
		}
	}

	return []byte(output)
}

func Html2Text(input io.Reader) string {
	return Xml2Text(input, []string{ "br", "p", "h1", "h2", "h3", "h4" }, []string{}, false)
}

var readabilityStopList map[string]bool = map[string]bool{ "and":true, "the":true, "a":true, "about":true, "above":true, "across":true, "after":true, "afterwards":true, "again":true, "against":true, "all":true, "almost":true, "alone":true,
	"along":true, "already":true, "also":true, "although":true, "always":true, "am":true, "among":true, "amongst":true, "amoungst":true, "amount":true, "an":true, "another":true, "any":true,
	"anyhow":true, "anyone":true, "anything":true, "anyway":true, "anywhere":true, "are":true, "around":true, "as":true, "at":true, "back":true, "be":true, "became":true, "because":true,
	"become":true, "becomes":true, "becoming":true, "been":true, "before":true, "beforehand":true, "behind":true, "being":true, "below":true, "beside":true, "besides":true, "between":true,
	"beyond":true, "both":true, "bottom":true, "but":true, "by":true, "can":true, "cannot":true, "cant":true, "co":true, "con":true, "could":true, "couldnt":true, "cry":true,
	"de":true, "describe":true, "detail":true, "do":true, "done":true, "down":true, "due":true, "during":true, "each":true, "eg":true, "eight":true, "either":true, "eleven":true, "else":true,
	"elsewhere":true, "empty":true, "enough":true, "etc":true, "even":true, "ever":true, "every":true, "everyone":true, "everything":true, "everywhere":true, "except":true, "few":true,
	"fifteen":true, "fify":true, "fill":true, "find":true, "fire":true, "first":true, "five":true, "for":true, "former":true, "formerly":true, "forty":true, "found":true, "four":true, "from":true,
	"front":true, "full":true, "further":true, "get":true, "give":true, "go":true, "had":true, "has":true, "hasnt":true, "have":true, "he":true, "hence":true, "her":true, "here":true, "hereafter":true,
	"hereby":true, "herein":true, "hereupon":true, "hers":true, "herself":true, "him":true, "himself":true, "his":true, "how":true, "however":true, "hundred":true, "ie":true, "if":true, "in":true,
	"inc":true, "indeed":true, "interest":true, "into":true, "is":true, "it":true, "its":true, "itself":true, "keep":true, "last":true, "latter":true, "latterly":true, "least":true, "less":true,
	"ltd":true, "made":true, "many":true, "may":true, "me":true, "meanwhile":true, "might":true, "mill":true, "mine":true, "more":true, "moreover":true, "most":true, "mostly":true, "move":true,
	"much":true, "must":true, "my":true, "myself":true, "name":true, "namely":true, "neither":true, "never":true, "nevertheless":true, "next":true, "nine":true, "no":true, "nobody":true,
	"none":true, "noone":true, "nor":true, "not":true, "nothing":true, "now":true, "nowhere":true, "of":true, "off":true, "often":true, "on":true, "once":true, "one":true, "only":true, "onto":true,
	"or":true, "other":true, "others":true, "otherwise":true, "our":true, "ours":true, "ourselves":true, "out":true, "over":true, "own":true, "part":true, "per":true, "perhaps":true,
	"please":true, "put":true, "rather":true, "re":true, "same":true, "see":true, "seem":true, "seemed":true, "seeming":true, "seems":true, "serious":true, "several":true, "she":true,
	"should":true, "show":true, "side":true, "since":true, "sincere":true, "six":true, "sixty":true, "so":true, "some":true, "somehow":true, "someone":true, "something":true, "sometime":true,
	"sometimes":true, "somewhere":true, "still":true, "such":true, "take":true, "ten":true, "than":true, "that":true, "their":true, "them":true, "themselves":true,
	"then":true, "thence":true, "there":true, "thereafter":true, "thereby":true, "therefore":true, "therein":true, "thereupon":true, "these":true, "they":true, "thickv":true, "thin":true,
	"third":true, "this":true, "those":true, "though":true, "three":true, "through":true, "throughout":true, "thru":true, "thus":true, "to":true, "together":true, "too":true, "top":true,
	"toward":true, "towards":true, "twelve":true, "twenty":true, "two":true, "un":true, "under":true, "until":true, "up":true, "upon":true, "us":true, "very":true, "via":true, "was":true, "we":true,
	"well":true, "were":true, "what":true, "whatever":true, "when":true, "whence":true, "whenever":true, "where":true, "whereafter":true, "whereas":true, "whereby":true, "wherein":true,
	"whereupon":true, "wherever":true, "whether":true, "which":true, "while":true, "whither":true, "who":true, "whoever":true, "whole":true, "whom":true, "whose":true, "why":true, "will":true,
	"with":true, "within":true, "without":true, "would":true, "yet":true, "you":true, "your":true, "youre":true, "yours":true, "yourself":true, "yourselves":true, "www":true, "com":true, "http":true }
<|endoftext|>


<commit_before>package dry

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

type wrappedResponseWriter struct {
	io.Writer
	http.ResponseWriter
}

func (wrapped wrappedResponseWriter) Write(data []byte) (int, error) {
	return wrapped.Writer.Write(data)
}

// HTTPCompressHandlerFunc wraps a http.HandlerFunc so that the response gets
// gzip or deflate compressed if the Accept-Encoding header of the request allows it.
func HTTPCompressHandlerFunc(handlerFunc http.HandlerFunc) http.HandlerFunc {
	return func(response http.ResponseWriter, request *http.Request) {
		NewHTTPCompressHandlerFromFunc(handlerFunc).ServeHTTP(response, request)
	}
}

// HTTPCompressHandler wraps a http.Handler so that the response gets
// gzip or deflate compressed if the Accept-Encoding header of the request allows it.
type HTTPCompressHandler struct {
	http.Handler
}

func NewHTTPCompressHandler(handler http.Handler) *HTTPCompressHandler {
	return &HTTPCompressHandler{handler}
}

func NewHTTPCompressHandlerFromFunc(handler http.HandlerFunc) *HTTPCompressHandler {
	return &HTTPCompressHandler{handler}
}

func (h *HTTPCompressHandler) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	accept := request.Header.Get("Accept-Encoding")
	if strings.Contains(accept, "gzip") {
		response.Header().Set("Content-Encoding", "gzip")
		writer := Gzip.GetWriter(response)
		defer Gzip.ReturnWriter(writer)
		response = wrappedResponseWriter{Writer: writer, ResponseWriter: response}
	} else if strings.Contains(accept, "deflate") {
		response.Header().Set("Content-Encoding", "deflate")
		writer := Deflate.GetWriter(response)
		defer Deflate.ReturnWriter(writer)
		response = wrappedResponseWriter{Writer: writer, ResponseWriter: response}
	}
	h.Handler.ServeHTTP(response, request)
}

// HTTPPostJSON marshalles data as JSON
// and sends it as HTTP POST request to url.
// If the response status code is not 200 OK,
// then the status is returned as an error.
func HTTPPostJSON(url string, data interface{}) error {
	b, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return err
	}
	response, err := http.Post(url, "application/json", bytes.NewBuffer(b))
	if err == nil && (response.StatusCode < 200 || response.StatusCode > 299) {
		err = errors.New(response.Status)
	}
	return err
}

// HTTPPostXML marshalles data as XML
// and sends it as HTTP POST request to url.
// If the response status code is not 200 OK,
// then the status is returned as an error.
func HTTPPostXML(url string, data interface{}) error {
	b, err := xml.MarshalIndent(data, "", "  ")
	if err != nil {
		return err
	}
	response, err := http.Post(url, "application/xml", bytes.NewBuffer(b))
	if err == nil && (response.StatusCode < 200 || response.StatusCode > 299) {
		err = errors.New(response.Status)
	}
	return err
}

// HTTPDelete performs a HTTP DELETE request
func HTTPDelete(url string) (statusCode int, statusText string, err error) {
	request, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return 0, "", err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, "", err
	}
	return response.StatusCode, response.Status, nil
}

// HTTPPostForm performs a HTTP POST request with data as application/x-www-form-urlencoded
func HTTPPostForm(url string, data url.Values) (statusCode int, statusText string, err error) {
	request, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return 0, "", err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, "", err
	}
	return response.StatusCode, response.Status, nil
}

// HTTPPutForm performs a HTTP PUT request with data as application/x-www-form-urlencoded
func HTTPPutForm(url string, data url.Values) (statusCode int, statusText string, err error) {
	request, err := http.NewRequest("PUT", url, strings.NewReader(data.Encode()))
	if err != nil {
		return 0, "", err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, "", err
	}
	return response.StatusCode, response.Status, nil
}

// HTTPRespondMarshalJSON marshals response as JSON to responseWriter, sets Content-Type to application/json
// and compresses the response if Content-Encoding from the request allows it.
func HTTPRespondMarshalJSON(response interface{}, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		if data, err = json.Marshal(response); err == nil {
			responseWriter.Header().Set("Content-Type", "application/json")
			_, err = responseWriter.Write(data)
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondMarshalIndentJSON marshals response as JSON to responseWriter, sets Content-Type to application/json
// and compresses the response if Content-Encoding from the request allows it.
// The JSON will be marshalled indented according to json.MarshalIndent
func HTTPRespondMarshalIndentJSON(response interface{}, prefix, indent string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		if data, err = json.MarshalIndent(response, prefix, indent); err == nil {
			responseWriter.Header().Set("Content-Type", "application/json")
			_, err = responseWriter.Write(data)
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondMarshalXML marshals response as XML to responseWriter, sets Content-Type to application/xml
// and compresses the response if Content-Encoding from the request allows it.
// If rootElement is not empty, then an additional root element with this name will be wrapped around the content.
func HTTPRespondMarshalXML(response interface{}, rootElement string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		if data, err = xml.Marshal(response); err == nil {
			responseWriter.Header().Set("Content-Type", "application/xml")
			if rootElement == "" {
				_, err = fmt.Fprintf(responseWriter, "%s%s", xml.Header, data)
			} else {
				_, err = fmt.Fprintf(responseWriter, "%s<%s>%s</%s>", xml.Header, rootElement, data, rootElement)
			}
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondMarshalIndentXML marshals response as XML to responseWriter, sets Content-Type to application/xml
// and compresses the response if Content-Encoding from the request allows it.
// The XML will be marshalled indented according to xml.MarshalIndent.
// If rootElement is not empty, then an additional root element with this name will be wrapped around the content.
func HTTPRespondMarshalIndentXML(response interface{}, rootElement string, prefix, indent string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		contentPrefix := prefix
		if rootElement != "" {
			contentPrefix += indent
		}
		if data, err = xml.MarshalIndent(response, contentPrefix, indent); err == nil {
			responseWriter.Header().Set("Content-Type", "application/xml")
			if rootElement == "" {
				_, err = fmt.Fprintf(responseWriter, "%s%s\n%s", prefix, xml.Header, data)
			} else {
				_, err = fmt.Fprintf(responseWriter, "%s%s%s<%s>\n%s\n%s</%s>", prefix, xml.Header, prefix, rootElement, data, prefix, rootElement)
			}
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondText sets Content-Type to text/plain
// and compresses the response if Content-Encoding from the request allows it.
func HTTPRespondText(response string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		responseWriter.Header().Set("Content-Type", "text/plain")
		_, err = responseWriter.Write([]byte(response))
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPUnmarshalRequestBodyJSON reads a http.Request body and unmarshals it as JSON to result.
func HTTPUnmarshalRequestBodyJSON(request *http.Request, result interface{}) error {
	body, err := ioutil.ReadAll(request.Body)
	if err != nil {
		return err
	}
	defer request.Body.Close()
	return json.Unmarshal(body, result)
}
<commit_msg>fix Body.Close() location in HTTPUnmarshalRequestBodyJSON<commit_after>package dry

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

type wrappedResponseWriter struct {
	io.Writer
	http.ResponseWriter
}

func (wrapped wrappedResponseWriter) Write(data []byte) (int, error) {
	return wrapped.Writer.Write(data)
}

// HTTPCompressHandlerFunc wraps a http.HandlerFunc so that the response gets
// gzip or deflate compressed if the Accept-Encoding header of the request allows it.
func HTTPCompressHandlerFunc(handlerFunc http.HandlerFunc) http.HandlerFunc {
	return func(response http.ResponseWriter, request *http.Request) {
		NewHTTPCompressHandlerFromFunc(handlerFunc).ServeHTTP(response, request)
	}
}

// HTTPCompressHandler wraps a http.Handler so that the response gets
// gzip or deflate compressed if the Accept-Encoding header of the request allows it.
type HTTPCompressHandler struct {
	http.Handler
}

func NewHTTPCompressHandler(handler http.Handler) *HTTPCompressHandler {
	return &HTTPCompressHandler{handler}
}

func NewHTTPCompressHandlerFromFunc(handler http.HandlerFunc) *HTTPCompressHandler {
	return &HTTPCompressHandler{handler}
}

func (h *HTTPCompressHandler) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	accept := request.Header.Get("Accept-Encoding")
	if strings.Contains(accept, "gzip") {
		response.Header().Set("Content-Encoding", "gzip")
		writer := Gzip.GetWriter(response)
		defer Gzip.ReturnWriter(writer)
		response = wrappedResponseWriter{Writer: writer, ResponseWriter: response}
	} else if strings.Contains(accept, "deflate") {
		response.Header().Set("Content-Encoding", "deflate")
		writer := Deflate.GetWriter(response)
		defer Deflate.ReturnWriter(writer)
		response = wrappedResponseWriter{Writer: writer, ResponseWriter: response}
	}
	h.Handler.ServeHTTP(response, request)
}

// HTTPPostJSON marshalles data as JSON
// and sends it as HTTP POST request to url.
// If the response status code is not 200 OK,
// then the status is returned as an error.
func HTTPPostJSON(url string, data interface{}) error {
	b, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return err
	}
	response, err := http.Post(url, "application/json", bytes.NewBuffer(b))
	if err == nil && (response.StatusCode < 200 || response.StatusCode > 299) {
		err = errors.New(response.Status)
	}
	return err
}

// HTTPPostXML marshalles data as XML
// and sends it as HTTP POST request to url.
// If the response status code is not 200 OK,
// then the status is returned as an error.
func HTTPPostXML(url string, data interface{}) error {
	b, err := xml.MarshalIndent(data, "", "  ")
	if err != nil {
		return err
	}
	response, err := http.Post(url, "application/xml", bytes.NewBuffer(b))
	if err == nil && (response.StatusCode < 200 || response.StatusCode > 299) {
		err = errors.New(response.Status)
	}
	return err
}

// HTTPDelete performs a HTTP DELETE request
func HTTPDelete(url string) (statusCode int, statusText string, err error) {
	request, err := http.NewRequest("DELETE", url, nil)
	if err != nil {
		return 0, "", err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, "", err
	}
	return response.StatusCode, response.Status, nil
}

// HTTPPostForm performs a HTTP POST request with data as application/x-www-form-urlencoded
func HTTPPostForm(url string, data url.Values) (statusCode int, statusText string, err error) {
	request, err := http.NewRequest("POST", url, strings.NewReader(data.Encode()))
	if err != nil {
		return 0, "", err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, "", err
	}
	return response.StatusCode, response.Status, nil
}

// HTTPPutForm performs a HTTP PUT request with data as application/x-www-form-urlencoded
func HTTPPutForm(url string, data url.Values) (statusCode int, statusText string, err error) {
	request, err := http.NewRequest("PUT", url, strings.NewReader(data.Encode()))
	if err != nil {
		return 0, "", err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, "", err
	}
	return response.StatusCode, response.Status, nil
}

// HTTPRespondMarshalJSON marshals response as JSON to responseWriter, sets Content-Type to application/json
// and compresses the response if Content-Encoding from the request allows it.
func HTTPRespondMarshalJSON(response interface{}, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		if data, err = json.Marshal(response); err == nil {
			responseWriter.Header().Set("Content-Type", "application/json")
			_, err = responseWriter.Write(data)
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondMarshalIndentJSON marshals response as JSON to responseWriter, sets Content-Type to application/json
// and compresses the response if Content-Encoding from the request allows it.
// The JSON will be marshalled indented according to json.MarshalIndent
func HTTPRespondMarshalIndentJSON(response interface{}, prefix, indent string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		if data, err = json.MarshalIndent(response, prefix, indent); err == nil {
			responseWriter.Header().Set("Content-Type", "application/json")
			_, err = responseWriter.Write(data)
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondMarshalXML marshals response as XML to responseWriter, sets Content-Type to application/xml
// and compresses the response if Content-Encoding from the request allows it.
// If rootElement is not empty, then an additional root element with this name will be wrapped around the content.
func HTTPRespondMarshalXML(response interface{}, rootElement string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		if data, err = xml.Marshal(response); err == nil {
			responseWriter.Header().Set("Content-Type", "application/xml")
			if rootElement == "" {
				_, err = fmt.Fprintf(responseWriter, "%s%s", xml.Header, data)
			} else {
				_, err = fmt.Fprintf(responseWriter, "%s<%s>%s</%s>", xml.Header, rootElement, data, rootElement)
			}
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondMarshalIndentXML marshals response as XML to responseWriter, sets Content-Type to application/xml
// and compresses the response if Content-Encoding from the request allows it.
// The XML will be marshalled indented according to xml.MarshalIndent.
// If rootElement is not empty, then an additional root element with this name will be wrapped around the content.
func HTTPRespondMarshalIndentXML(response interface{}, rootElement string, prefix, indent string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		var data []byte
		contentPrefix := prefix
		if rootElement != "" {
			contentPrefix += indent
		}
		if data, err = xml.MarshalIndent(response, contentPrefix, indent); err == nil {
			responseWriter.Header().Set("Content-Type", "application/xml")
			if rootElement == "" {
				_, err = fmt.Fprintf(responseWriter, "%s%s\n%s", prefix, xml.Header, data)
			} else {
				_, err = fmt.Fprintf(responseWriter, "%s%s%s<%s>\n%s\n%s</%s>", prefix, xml.Header, prefix, rootElement, data, prefix, rootElement)
			}
		}
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPRespondText sets Content-Type to text/plain
// and compresses the response if Content-Encoding from the request allows it.
func HTTPRespondText(response string, responseWriter http.ResponseWriter, request *http.Request) (err error) {
	NewHTTPCompressHandlerFromFunc(func(responseWriter http.ResponseWriter, request *http.Request) {
		responseWriter.Header().Set("Content-Type", "text/plain")
		_, err = responseWriter.Write([]byte(response))
	}).ServeHTTP(responseWriter, request)
	return err
}

// HTTPUnmarshalRequestBodyJSON reads a http.Request body and unmarshals it as JSON to result.
func HTTPUnmarshalRequestBodyJSON(request *http.Request, result interface{}) error {
	defer request.Body.Close()
	body, err := ioutil.ReadAll(request.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(body, result)
}
<|endoftext|>


<commit_before>package iocap

import (
	"net/http"
)

// LimitedHTTPHandler is a wrapper over a normal http.Handler, allowing the
// rate to be controlled while sending data back to clients.
type LimitedHTTPHandler struct {
	h    http.Handler
	opts RateOpts
}

// LimitHTTPHandler creates a new rate limited HTTP handler wrapper.
func LimitHTTPHandler(h http.Handler, opts RateOpts) *LimitedHTTPHandler {
	return &LimitedHTTPHandler{
		h:    h,
		opts: opts,
	}
}

// ServeHTTP implements the http.Handler interface, writing responses using
// a rate limited response writer.
func (h *LimitedHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.h.ServeHTTP(LimitResponseWriter(w, h.opts), r)
}

// LimitedResponseWriter wraps an http.ResponseWriter in a rate limited
// writer, effectively throttling throughput from the HTTP server to
// all of its clients.
type LimitedResponseWriter struct {
	writer *Writer
	http.ResponseWriter
}

// LimitResponseWriter creates a new rate limited, wrapped response writer.
func LimitResponseWriter(w http.ResponseWriter, opts RateOpts) *LimitedResponseWriter {
	return &LimitedResponseWriter{
		writer:         NewWriter(w, opts),
		ResponseWriter: w,
	}
}

// Write implements part of the http.ResponseWriter interface, calling the
// underlying rate limited writer instead of directly writing out bytes.
func (w *LimitedResponseWriter) Write(p []byte) (int, error) {
	return w.writer.Write(p)
}
<commit_msg>Cleanup<commit_after>package iocap

import "net/http"

// LimitedHTTPHandler is a wrapper over a normal http.Handler, allowing the
// rate to be controlled while sending data back to clients.
type LimitedHTTPHandler struct {
	h    http.Handler
	opts RateOpts
}

// LimitHTTPHandler creates a new rate limited HTTP handler wrapper.
func LimitHTTPHandler(h http.Handler, opts RateOpts) *LimitedHTTPHandler {
	return &LimitedHTTPHandler{
		h:    h,
		opts: opts,
	}
}

// ServeHTTP implements the http.Handler interface, writing responses using
// a rate limited response writer.
func (h *LimitedHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.h.ServeHTTP(LimitResponseWriter(w, h.opts), r)
}

// LimitedResponseWriter wraps an http.ResponseWriter in a rate limited
// writer, effectively throttling throughput from the HTTP server to
// all of its clients.
type LimitedResponseWriter struct {
	writer *Writer
	http.ResponseWriter
}

// LimitResponseWriter creates a new rate limited, wrapped response writer.
func LimitResponseWriter(w http.ResponseWriter, opts RateOpts) *LimitedResponseWriter {
	return &LimitedResponseWriter{
		writer:         NewWriter(w, opts),
		ResponseWriter: w,
	}
}

// Write implements part of the http.ResponseWriter interface, calling the
// underlying rate limited writer instead of directly writing out bytes.
func (w *LimitedResponseWriter) Write(p []byte) (int, error) {
	return w.writer.Write(p)
}
<|endoftext|>


<commit_before>package reverseproxy

import (
	"io"
	"net/http"
)

var transport *http.Transport = &http.Transport{DisableKeepAlives: true}

// ProxyHTTP proxies an HTTP request to a given destination host.
// This will not handle WebSockets intelligently.
func ProxyHTTP(w http.ResponseWriter, r *http.Request, host string) {
	proxyHTTP(w, r, []string{host}, []int{0})
}

func proxyHTTP(w http.ResponseWriter, r *http.Request, hosts []string,
	indices []int) {
	var res *http.Response
	var err error
	for _, i := range indices {
		host := hosts[i]

		// Generate the request for the proxy destination.
		targetURL := *r.URL
		targetURL.Host = host
		targetURL.Scheme = "http"
		var req *http.Request
		req, err = http.NewRequest(r.Method, targetURL.String(), r.Body)
		if err != nil {
			continue
		}
		req.Header = requestHeaders(r, host, false)
		req.Host = r.Host

		// Send the request
		res, err = transport.RoundTrip(req)
		// If an error occurs the request's body may have been read and trying a
		// new host would be pointless. However, I do not currently check for
		// this.
		if err != nil {
			continue
		}
		break
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}

	// Write the response
	respHead := responseHeaders(res, r.Host, r.URL.Scheme, false)
	for header, value := range respHead {
		w.Header()[header] = value
	}
	w.WriteHeader(res.StatusCode)
	io.Copy(w, res.Body)
	res.Body.Close()
}
<commit_msg>flush chunked/streaming HTTP responses<commit_after>package reverseproxy

import (
	"io"
	"net/http"
)

var transport *http.Transport = &http.Transport{DisableKeepAlives: true}

// ProxyHTTP proxies an HTTP request to a given destination host.
// This will not handle WebSockets intelligently.
func ProxyHTTP(w http.ResponseWriter, r *http.Request, host string) {
	proxyHTTP(w, r, []string{host}, []int{0})
}

func proxyHTTP(w http.ResponseWriter, r *http.Request, hosts []string,
	indices []int) {
	var res *http.Response
	var err error
	for _, i := range indices {
		host := hosts[i]

		// Generate the request for the proxy destination.
		targetURL := *r.URL
		targetURL.Host = host
		targetURL.Scheme = "http"
		var req *http.Request
		req, err = http.NewRequest(r.Method, targetURL.String(), r.Body)
		if err != nil {
			continue
		}
		req.Header = requestHeaders(r, host, false)
		req.Host = r.Host

		// Send the request
		res, err = transport.RoundTrip(req)
		// If an error occurs the request's body may have been read and trying a
		// new host would be pointless. However, I do not currently check for
		// this.
		if err != nil {
			continue
		}
		break
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}

	// Write the response
	respHead := responseHeaders(res, r.Host, r.URL.Scheme, false)
	for header, value := range respHead {
		w.Header()[header] = value
	}
	w.WriteHeader(res.StatusCode)
	if respHead.Get("X-Content-Type-Options") == "nosniff" {
		// This header tends to be set for chunked/streaming responses.
		io.Copy(&httpFlushWriter{w: w, f: w.(http.Flusher)}, res.Body)
	} else {
		io.Copy(w, res.Body)
	}
	res.Body.Close()
}

type httpFlushWriter struct {
	w io.Writer
	f http.Flusher
}

func (h *httpFlushWriter) Write(buf []byte) (int, error) {
	n, err := h.w.Write(buf)
	h.f.Flush()
	return n, err
}
<|endoftext|>


<commit_before>package ranger

import (
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"
)

// HTTPRanger is a RangeFetcher that uses the HTTP Range: header to fetch blocks.
//
// HTTPRanger first makes a HEAD request and then between 0 and Length()/BlockSize GET requests, attempting
// whenever possible to optimize for a lower number of requests.
//
// No network requests are made until the first I/O-related function call.
type HTTPRanger struct {
	URL                *url.URL
	etag, lastModified string
	length             int64
	client             http.Client
	blockSize          int
}

func statusCodeError(status int) error {
	return fmt.Errorf("unexpected response (status %d)", status)
}

func statusIsAcceptable(status int) bool {
	return status >= 200 && status < 300
}

// Initialize implements the Initialize function from the RangeFetcher interface.
// It performs a HEAD request to retrieve the required information from the server.
func (r *HTTPRanger) Initialize(bs int) error {
	resp, err := r.client.Head(r.URL.String())
	if err != nil {
		return err
	}

	if !statusIsAcceptable(resp.StatusCode) {
		return statusCodeError(resp.StatusCode)
	}

	if !strings.Contains(resp.Header.Get("Accept-Ranges"), "bytes") {
		return errors.New(r.URL.Host + " does not support byte-ranged requests.")
	}

	r.blockSize = bs
	r.etag = resp.Header.Get("ETag")
	r.lastModified = resp.Header.Get("Last-Modified")
	r.length = resp.ContentLength
	return nil
}

// Length returns the length, in bytes, of the ranged-over file.
func (r *HTTPRanger) Length() int64 {
	return r.length
}

// FetchBlocks requests blocks from the HTTP server.
func (r *HTTPRanger) FetchBlocks(ranges []BlockByteRange) ([]Block, error) {
	blox := make([]Block, len(ranges))
	if len(ranges) > 0 {
		rs := make([]string, len(ranges))
		for i, rng := range ranges {
			rs[i] = fmt.Sprintf("%d-%d", rng.Start, rng.End)
		}
		rangeString := strings.Join(rs, ",")

		req, _ := http.NewRequest("GET", r.URL.String(), nil)
		req.Header.Set("Range", fmt.Sprintf("bytes=%s", rangeString))
		if r.etag != "" {
			req.Header.Set("If-Range", r.etag)
		} else if r.lastModified != "" {
			req.Header.Set("If-Range", r.lastModified)
		}

		resp, err := r.client.Do(req)
		if err != nil {
			return nil, err
		}

		switch resp.StatusCode {
		case http.StatusPreconditionFailed:
			return nil, errors.New("ranged request failed; document may have mutated")
		case http.StatusNotFound:
			return nil, errors.New("ranged request failed; document may have disappeared")
		default:
			if !statusIsAcceptable(resp.StatusCode) {
				return nil, statusCodeError(resp.StatusCode)
			}
		}

		defer resp.Body.Close()

		typ, params, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
		if err != nil {
			return nil, err
		}

		if typ == "multipart/byteranges" {
			multipart := multipart.NewReader(resp.Body, params["boundary"])
			i := 0
			for {
				if part, err := multipart.NextPart(); err == nil {
					rng := ranges[i]
					bn := rng.Number
					blocklen := (rng.End - rng.Start) + 1
					blox[i] = Block{Number: bn, Data: make([]byte, blocklen)}
					io.ReadFull(part, blox[i].Data)
					i++
				} else {
					if err == io.EOF {
						break
					}
					return nil, err
				}
			}
		} else {
			bn := 0
			if resp.StatusCode == http.StatusPartialContent {
				// If we've received 206 Partial Content but no multipart parts,
				// we received a contiguous section starting at the first requested block.
				bn = ranges[0].Number
			}
			body := make([]byte, resp.ContentLength)
			io.ReadFull(resp.Body, body)
			blox = blox[0:0]
			remaining := resp.ContentLength
			ncopied := int64(0)
			for remaining > 0 {
				bs := int64(r.blockSize)
				if bs > remaining {
					bs = remaining
				}

				blk := Block{bn, make([]byte, bs)}
				bodySlice := body[ncopied : ncopied+bs]
				copy(blk.Data, bodySlice)
				blox = append(blox, blk)

				bn++
				ncopied += bs
				remaining -= bs
			}
		}
	}
	return blox, nil
}
<commit_msg>Make Client an exported field in HTTPRanger<commit_after>package ranger

import (
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"strings"
)

// HTTPClient is an interface describing the methods required from net/http.Client
type HTTPClient interface {
	Do(*http.Request) (*http.Response, error)
	Get(string) (*http.Response, error)
	Head(string) (*http.Response, error)
}

// HTTPRanger is a RangeFetcher that uses the HTTP Range: header to fetch blocks.
//
// HTTPRanger first makes a HEAD request and then between 0 and Length()/BlockSize GET requests, attempting
// whenever possible to optimize for a lower number of requests.
//
// No network requests are made until the first I/O-related function call.
type HTTPRanger struct {
	URL    *url.URL
	Client HTTPClient

	etag, lastModified string
	length             int64
	blockSize          int
}

func statusCodeError(status int) error {
	return fmt.Errorf("unexpected response (status %d)", status)
}

func statusIsAcceptable(status int) bool {
	return status >= 200 && status < 300
}

// Initialize implements the Initialize function from the RangeFetcher interface.
// It performs a HEAD request to retrieve the required information from the server.
func (r *HTTPRanger) Initialize(bs int) error {
	if r.Client == nil {
		r.Client = &http.Client{}
	}

	resp, err := r.Client.Head(r.URL.String())
	if err != nil {
		return err
	}

	if !statusIsAcceptable(resp.StatusCode) {
		return statusCodeError(resp.StatusCode)
	}

	if !strings.Contains(resp.Header.Get("Accept-Ranges"), "bytes") {
		return errors.New(r.URL.Host + " does not support byte-ranged requests.")
	}

	r.blockSize = bs
	r.etag = resp.Header.Get("ETag")
	r.lastModified = resp.Header.Get("Last-Modified")
	r.length = resp.ContentLength
	return nil
}

// Length returns the length, in bytes, of the ranged-over file.
func (r *HTTPRanger) Length() int64 {
	return r.length
}

// FetchBlocks requests blocks from the HTTP server.
func (r *HTTPRanger) FetchBlocks(ranges []BlockByteRange) ([]Block, error) {
	blox := make([]Block, len(ranges))
	if len(ranges) > 0 {
		rs := make([]string, len(ranges))
		for i, rng := range ranges {
			rs[i] = fmt.Sprintf("%d-%d", rng.Start, rng.End)
		}
		rangeString := strings.Join(rs, ",")

		req, _ := http.NewRequest("GET", r.URL.String(), nil)
		req.Header.Set("Range", fmt.Sprintf("bytes=%s", rangeString))
		if r.etag != "" {
			req.Header.Set("If-Range", r.etag)
		} else if r.lastModified != "" {
			req.Header.Set("If-Range", r.lastModified)
		}

		resp, err := r.Client.Do(req)
		if err != nil {
			return nil, err
		}

		switch resp.StatusCode {
		case http.StatusPreconditionFailed:
			return nil, errors.New("ranged request failed; document may have mutated")
		case http.StatusNotFound:
			return nil, errors.New("ranged request failed; document may have disappeared")
		default:
			if !statusIsAcceptable(resp.StatusCode) {
				return nil, statusCodeError(resp.StatusCode)
			}
		}

		defer resp.Body.Close()

		typ, params, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
		if err != nil {
			return nil, err
		}

		if typ == "multipart/byteranges" {
			multipart := multipart.NewReader(resp.Body, params["boundary"])
			i := 0
			for {
				if part, err := multipart.NextPart(); err == nil {
					rng := ranges[i]
					bn := rng.Number
					blocklen := (rng.End - rng.Start) + 1
					blox[i] = Block{Number: bn, Data: make([]byte, blocklen)}
					io.ReadFull(part, blox[i].Data)
					i++
				} else {
					if err == io.EOF {
						break
					}
					return nil, err
				}
			}
		} else {
			bn := 0
			if resp.StatusCode == http.StatusPartialContent {
				// If we've received 206 Partial Content but no multipart parts,
				// we received a contiguous section starting at the first requested block.
				bn = ranges[0].Number
			}
			body := make([]byte, resp.ContentLength)
			io.ReadFull(resp.Body, body)
			blox = blox[0:0]
			remaining := resp.ContentLength
			ncopied := int64(0)
			for remaining > 0 {
				bs := int64(r.blockSize)
				if bs > remaining {
					bs = remaining
				}

				blk := Block{bn, make([]byte, bs)}
				bodySlice := body[ncopied : ncopied+bs]
				copy(blk.Data, bodySlice)
				blox = append(blox, blk)

				bn++
				ncopied += bs
				remaining -= bs
			}
		}
	}
	return blox, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"sort"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/newrelic/sidecar/catalog"
	"github.com/newrelic/sidecar/output"
	"github.com/newrelic/sidecar/service"
	"github.com/nitro/memberlist"
)

type ApiServer struct {
	Name         string
	LastUpdated  time.Time
	ServiceCount int
}

type ApiServices struct {
	Services       map[string][]*service.Service
	ClusterMembers map[string]*ApiServer
	ClusterName    string
}

func makeHandler(fn func(http.ResponseWriter, *http.Request,
	*memberlist.Memberlist, *catalog.ServicesState),
	list *memberlist.Memberlist, state *catalog.ServicesState) http.HandlerFunc {

	return func(response http.ResponseWriter, req *http.Request) {
		fn(response, req, list, state)
	}
}

func watchHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "application/json")

	lastChange := time.Unix(0, 0)
	var jsonBytes []byte
	var err error

	for {
		var changed bool

		func() { // Wrap critical section
			state.RLock()
			defer state.RUnlock()

			if state.LastChanged.After(lastChange) {
				lastChange = state.LastChanged
				jsonBytes, err = json.Marshal(state.ByService())
				if err != nil {
					log.Errorf("Error marshaling state in watchHandler: %s", err.Error())
					return
				}
				changed = true // Trigger sending new encoding
			}
		}()

		if changed {
			// In order to flush immediately, we have to cast to a Flusher.
			// The normal HTTP library supports this but not all do, so we
			// check just in case.
			if f, ok := response.(http.Flusher); ok {
				f.Flush()
			}
			response.Write(jsonBytes)
		}

		time.Sleep(250 * time.Millisecond)
	}
}

func servicesHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	params := mux.Vars(req)
	defer req.Body.Close()

	if params["extension"] == ".json" {
		response.Header().Set("Content-Type", "application/json")
		response.Header().Set("Access-Control-Allow-Origin", "*")
		response.Header().Set("Access-Control-Allow-Methods", "GET")

		listMembers := list.Members()
		sort.Sort(listByName(listMembers))
		members := make(map[string]*ApiServer, len(listMembers))

		var jsonBytes []byte
		var err error

		func() { // Wrap critical section
			state.RLock()
			defer state.RUnlock()

			for _, member := range listMembers {
				if state.HasServer(member.Name) {
					members[member.Name] = &ApiServer{
						Name:         member.Name,
						LastUpdated:  state.Servers[member.Name].LastUpdated,
						ServiceCount: len(state.Servers[member.Name].Services),
					}
				} else {
					members[member.Name] = &ApiServer{
						Name:         member.Name,
						LastUpdated:  time.Unix(0, 0),
						ServiceCount: 0,
					}
				}
			}

			result := ApiServices{
				Services:       state.ByService(),
				ClusterMembers: members,
				ClusterName:    list.ClusterName(),
			}

			jsonBytes, err = json.MarshalIndent(&result, "", "  ")
		}()

		if err != nil {
			log.Errorf("Error marshaling state in servicesHandler: %s", err.Error())
		}

		response.Write(jsonBytes)
		return
	}
}

func serversHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "text/html")
	state.RLock()
	defer state.RUnlock()

	response.Write(
		[]byte(`

 			<meta http-equiv="refresh" content="4">

` + state.Format(list) + "
"))
}

func stateHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	params := mux.Vars(req)
	defer req.Body.Close()

	state.RLock()
	defer state.RUnlock()

	if params["extension"] == ".json" {
		response.Header().Set("Content-Type", "application/json")
		response.Header().Set("Access-Control-Allow-Origin", "*")
		response.Header().Set("Access-Control-Allow-Methods", "GET")
		response.Write(state.Encode())
		return
	}
}

func optionsHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	response.Header().Set("Access-Control-Allow-Origin", "*")
	response.Header().Set("Access-Control-Allow-Methods", "GET")
	return
}

func statusStr(status int) string {
	switch status {
	case 0:
		return "Alive"
	case 1:
		return "Tombstone"
	case 2:
		return "Unhealthy"
	default:
		return "Unknown"
	}
}

func portsStr(svcPorts []service.Port) string {
	var ports []string

	for _, port := range svcPorts {
		if port.ServicePort != 0 {
			ports = append(ports, fmt.Sprintf("%v->%v", port.ServicePort, port.Port))
		} else {
			ports = append(ports, fmt.Sprintf("%v", port.Port))
		}
	}

	return strings.Join(ports, ", ")
}

type listByName []*memberlist.Node

func (a listByName) Len() int           { return len(a) }
func (a listByName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a listByName) Less(i, j int) bool { return a[i].Name < a[j].Name }

type Member struct {
	Node    *memberlist.Node
	Updated time.Time
}

func lineWrapMembers(cols int, fields []*Member) [][]*Member {
	if len(fields) < cols {
		return [][]*Member{fields}
	}

	retval := make([][]*Member, len(fields)/cols+1)
	for i := 0; i < len(fields); i++ {
		row := i / cols
		retval[row] = append(retval[row], fields[i])
	}

	return retval
}

func viewHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	timeAgo := func(when time.Time) string { return output.TimeAgo(when, time.Now().UTC()) }

	funcMap := template.FuncMap{
		"statusStr":   statusStr,
		"timeAgo":     timeAgo,
		"portsStr":    portsStr,
		"clusterName": func() string { return list.ClusterName() },
	}

	t, err := template.New("services").Funcs(funcMap).ParseFiles("views/services.html")
	if err != nil {
		log.Errorf("Error parsing template: %s", err.Error())
	}

	members := list.Members()
	sort.Sort(listByName(members))

	state.RLock()
	defer state.RUnlock()

	compiledMembers := make([]*Member, len(members))
	for i, member := range members {
		if state.HasServer(member.Name) {
			compiledMembers[i] = &Member{member, state.Servers[member.Name].LastUpdated}
		} else {
			compiledMembers[i] = &Member{Node: member}
			log.Debug("No updated time for " + member.Name)
		}
	}

	wrappedMembers := lineWrapMembers(5, compiledMembers)

	viewData := struct {
		Services map[string][]*service.Service
		Members  [][]*Member
	}{
		Services: state.ByService(),
		Members:  wrappedMembers,
	}

	t.ExecuteTemplate(response, "services.html", viewData)
}

func uiRedirectHandler(response http.ResponseWriter, req *http.Request) {
	http.Redirect(response, req, "/ui/", 301)
}

func serveHttp(list *memberlist.Memberlist, state *catalog.ServicesState) {
	router := mux.NewRouter()

	router.HandleFunc("/", uiRedirectHandler).Methods("GET")

	router.HandleFunc(
		"/services{extension}", makeHandler(servicesHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/servers", makeHandler(serversHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/services", makeHandler(viewHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/state{extension}", makeHandler(stateHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/watch", makeHandler(watchHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/{path}", makeHandler(optionsHandler, list, state),
	).Methods("OPTIONS")

	staticFs := http.FileServer(http.Dir("views/static"))
	router.PathPrefix("/static").Handler(http.StripPrefix("/static", staticFs))

	uiFs := http.FileServer(http.Dir("ui/app"))
	router.PathPrefix("/ui").Handler(http.StripPrefix("/ui", uiFs))

	http.Handle("/", router)

	err := http.ListenAndServe("0.0.0.0:7777", nil)
	exitWithError(err, "Can't start HTTP server")
}
<commit_msg>pprof on by default. Performance is fine with it.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	_ "net/http/pprof"
	"sort"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/newrelic/sidecar/catalog"
	"github.com/newrelic/sidecar/output"
	"github.com/newrelic/sidecar/service"
	"github.com/nitro/memberlist"
)

type ApiServer struct {
	Name         string
	LastUpdated  time.Time
	ServiceCount int
}

type ApiServices struct {
	Services       map[string][]*service.Service
	ClusterMembers map[string]*ApiServer
	ClusterName    string
}

func makeHandler(fn func(http.ResponseWriter, *http.Request,
	*memberlist.Memberlist, *catalog.ServicesState),
	list *memberlist.Memberlist, state *catalog.ServicesState) http.HandlerFunc {

	return func(response http.ResponseWriter, req *http.Request) {
		fn(response, req, list, state)
	}
}

func watchHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "application/json")

	lastChange := time.Unix(0, 0)
	var jsonBytes []byte
	var err error

	for {
		var changed bool

		func() { // Wrap critical section
			state.RLock()
			defer state.RUnlock()

			if state.LastChanged.After(lastChange) {
				lastChange = state.LastChanged
				jsonBytes, err = json.Marshal(state.ByService())
				if err != nil {
					log.Errorf("Error marshaling state in watchHandler: %s", err.Error())
					return
				}
				changed = true // Trigger sending new encoding
			}
		}()

		if changed {
			// In order to flush immediately, we have to cast to a Flusher.
			// The normal HTTP library supports this but not all do, so we
			// check just in case.
			if f, ok := response.(http.Flusher); ok {
				f.Flush()
			}
			response.Write(jsonBytes)
		}

		time.Sleep(250 * time.Millisecond)
	}
}

func servicesHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	params := mux.Vars(req)
	defer req.Body.Close()

	if params["extension"] == ".json" {
		response.Header().Set("Content-Type", "application/json")
		response.Header().Set("Access-Control-Allow-Origin", "*")
		response.Header().Set("Access-Control-Allow-Methods", "GET")

		listMembers := list.Members()
		sort.Sort(listByName(listMembers))
		members := make(map[string]*ApiServer, len(listMembers))

		var jsonBytes []byte
		var err error

		func() { // Wrap critical section
			state.RLock()
			defer state.RUnlock()

			for _, member := range listMembers {
				if state.HasServer(member.Name) {
					members[member.Name] = &ApiServer{
						Name:         member.Name,
						LastUpdated:  state.Servers[member.Name].LastUpdated,
						ServiceCount: len(state.Servers[member.Name].Services),
					}
				} else {
					members[member.Name] = &ApiServer{
						Name:         member.Name,
						LastUpdated:  time.Unix(0, 0),
						ServiceCount: 0,
					}
				}
			}

			result := ApiServices{
				Services:       state.ByService(),
				ClusterMembers: members,
				ClusterName:    list.ClusterName(),
			}

			jsonBytes, err = json.MarshalIndent(&result, "", "  ")
		}()

		if err != nil {
			log.Errorf("Error marshaling state in servicesHandler: %s", err.Error())
		}

		response.Write(jsonBytes)
		return
	}
}

func serversHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "text/html")
	state.RLock()
	defer state.RUnlock()

	response.Write(
		[]byte(`

 			<meta http-equiv="refresh" content="4">

` + state.Format(list) + "
"))
}

func stateHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	params := mux.Vars(req)
	defer req.Body.Close()

	state.RLock()
	defer state.RUnlock()

	if params["extension"] == ".json" {
		response.Header().Set("Content-Type", "application/json")
		response.Header().Set("Access-Control-Allow-Origin", "*")
		response.Header().Set("Access-Control-Allow-Methods", "GET")
		response.Write(state.Encode())
		return
	}
}

func optionsHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	response.Header().Set("Access-Control-Allow-Origin", "*")
	response.Header().Set("Access-Control-Allow-Methods", "GET")
	return
}

func statusStr(status int) string {
	switch status {
	case 0:
		return "Alive"
	case 1:
		return "Tombstone"
	case 2:
		return "Unhealthy"
	default:
		return "Unknown"
	}
}

func portsStr(svcPorts []service.Port) string {
	var ports []string

	for _, port := range svcPorts {
		if port.ServicePort != 0 {
			ports = append(ports, fmt.Sprintf("%v->%v", port.ServicePort, port.Port))
		} else {
			ports = append(ports, fmt.Sprintf("%v", port.Port))
		}
	}

	return strings.Join(ports, ", ")
}

type listByName []*memberlist.Node

func (a listByName) Len() int           { return len(a) }
func (a listByName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a listByName) Less(i, j int) bool { return a[i].Name < a[j].Name }

type Member struct {
	Node    *memberlist.Node
	Updated time.Time
}

func lineWrapMembers(cols int, fields []*Member) [][]*Member {
	if len(fields) < cols {
		return [][]*Member{fields}
	}

	retval := make([][]*Member, len(fields)/cols+1)
	for i := 0; i < len(fields); i++ {
		row := i / cols
		retval[row] = append(retval[row], fields[i])
	}

	return retval
}

func viewHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	timeAgo := func(when time.Time) string { return output.TimeAgo(when, time.Now().UTC()) }

	funcMap := template.FuncMap{
		"statusStr":   statusStr,
		"timeAgo":     timeAgo,
		"portsStr":    portsStr,
		"clusterName": func() string { return list.ClusterName() },
	}

	t, err := template.New("services").Funcs(funcMap).ParseFiles("views/services.html")
	if err != nil {
		log.Errorf("Error parsing template: %s", err.Error())
	}

	members := list.Members()
	sort.Sort(listByName(members))

	state.RLock()
	defer state.RUnlock()

	compiledMembers := make([]*Member, len(members))
	for i, member := range members {
		if state.HasServer(member.Name) {
			compiledMembers[i] = &Member{member, state.Servers[member.Name].LastUpdated}
		} else {
			compiledMembers[i] = &Member{Node: member}
			log.Debug("No updated time for " + member.Name)
		}
	}

	wrappedMembers := lineWrapMembers(5, compiledMembers)

	viewData := struct {
		Services map[string][]*service.Service
		Members  [][]*Member
	}{
		Services: state.ByService(),
		Members:  wrappedMembers,
	}

	t.ExecuteTemplate(response, "services.html", viewData)
}

func uiRedirectHandler(response http.ResponseWriter, req *http.Request) {
	http.Redirect(response, req, "/ui/", 301)
}

func serveHttp(list *memberlist.Memberlist, state *catalog.ServicesState) {
	router := mux.NewRouter()

	router.HandleFunc("/", uiRedirectHandler).Methods("GET")

	router.HandleFunc(
		"/services{extension}", makeHandler(servicesHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/servers", makeHandler(serversHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/services", makeHandler(viewHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/state{extension}", makeHandler(stateHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/watch", makeHandler(watchHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/{path}", makeHandler(optionsHandler, list, state),
	).Methods("OPTIONS")

	staticFs := http.FileServer(http.Dir("views/static"))
	router.PathPrefix("/static").Handler(http.StripPrefix("/static", staticFs))

	uiFs := http.FileServer(http.Dir("ui/app"))
	router.PathPrefix("/ui").Handler(http.StripPrefix("/ui", uiFs))

	http.Handle("/", router)

	err := http.ListenAndServe("0.0.0.0:7777", nil)
	exitWithError(err, "Can't start HTTP server")
}
<|endoftext|>


<commit_before>package gost

import (
	"bufio"
	"crypto/tls"
	"encoding/base64"
	"github.com/golang/glog"
	"golang.org/x/net/http2"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	//"strings"
	"errors"
	"time"
)

type HttpServer struct {
	conn net.Conn
	Base *ProxyServer
}

func NewHttpServer(conn net.Conn, base *ProxyServer) *HttpServer {
	return &HttpServer{
		conn: conn,
		Base: base,
	}
}

// Default HTTP server handler
func (s *HttpServer) HandleRequest(req *http.Request) {
	glog.V(LINFO).Infof("[http] %s %s - %s %s", req.Method, s.conn.RemoteAddr(), req.Host, req.Proto)

	if glog.V(LDEBUG) {
		dump, _ := httputil.DumpRequest(req, false)
		glog.Infoln(string(dump))
	}

	if req.Method == "PRI" && req.ProtoMajor == 2 {
		glog.V(LWARNING).Infof("[http] %s <- %s : Not an HTTP2 server", s.conn.RemoteAddr(), req.Host)
		resp := "HTTP/1.1 400 Bad Request\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n"
		s.conn.Write([]byte(resp))
		return
	}

	valid := false
	u, p, _ := basicProxyAuth(req.Header.Get("Proxy-Authorization"))
	glog.V(LINFO).Infoln(u, p)
	for _, user := range s.Base.Node.Users {
		username := user.Username()
		password, _ := user.Password()
		if (u == username && p == password) ||
			(u == username && password == "") ||
			(username == "" && p == password) {
			valid = true
			break
		}
	}

	if len(s.Base.Node.Users) > 0 && !valid {
		glog.V(LWARNING).Infof("[http] %s <- %s : proxy authentication required", s.conn.RemoteAddr(), req.Host)
		resp := "HTTP/1.1 407 Proxy Authentication Required\r\n" +
			"Proxy-Authenticate: Basic realm=\"gost\"\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n"
		s.conn.Write([]byte(resp))
		return
	}

	req.Header.Del("Proxy-Authorization")

	// forward http request
	lastNode := s.Base.Chain.lastNode
	if lastNode != nil && (lastNode.Protocol == "http" || lastNode.Protocol == "") {
		s.forwardRequest(req)
		return
	}

	c, err := s.Base.Chain.Dial(req.Host)
	if err != nil {
		glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), req.Host, err)

		b := []byte("HTTP/1.1 503 Service unavailable\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
		glog.V(LDEBUG).Infof("[http] %s <- %s\n%s", s.conn.RemoteAddr(), req.Host, string(b))
		s.conn.Write(b)
		return
	}
	defer c.Close()

	if req.Method == http.MethodConnect {
		b := []byte("HTTP/1.1 200 Connection established\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
		glog.V(LDEBUG).Infof("[http] %s <- %s\n%s", s.conn.RemoteAddr(), req.Host, string(b))
		s.conn.Write(b)
	} else {
		req.Header.Del("Proxy-Connection")
		req.Header.Set("Connection", "Keep-Alive")

		if err = req.Write(c); err != nil {
			glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), req.Host, err)
			return
		}
	}

	glog.V(LINFO).Infof("[http] %s <-> %s", s.conn.RemoteAddr(), req.Host)
	s.Base.transport(s.conn, c)
	glog.V(LINFO).Infof("[http] %s >-< %s", s.conn.RemoteAddr(), req.Host)
}

func (s *HttpServer) forwardRequest(req *http.Request) {
	last := s.Base.Chain.lastNode
	if last == nil {
		return
	}
	cc, err := s.Base.Chain.GetConn()
	if err != nil {
		glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), last.Addr, err)

		b := []byte("HTTP/1.1 503 Service unavailable\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
		glog.V(LDEBUG).Infof("[http] %s <- %s\n%s", s.conn.RemoteAddr(), last.Addr, string(b))
		s.conn.Write(b)
		return
	}
	defer cc.Close()

	if len(last.Users) > 0 {
		user := last.Users[0]
		s := user.String()
		if _, set := user.Password(); !set {
			s += ":"
		}
		req.Header.Set("Proxy-Authorization",
			"Basic "+base64.StdEncoding.EncodeToString([]byte(s)))
	}

	cc.SetWriteDeadline(time.Now().Add(WriteTimeout))
	if err = req.WriteProxy(cc); err != nil {
		glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), req.Host, err)
		return
	}
	cc.SetWriteDeadline(time.Time{})

	glog.V(LINFO).Infof("[http] %s <-> %s", s.conn.RemoteAddr(), req.Host)
	s.Base.transport(s.conn, cc)
	glog.V(LINFO).Infof("[http] %s >-< %s", s.conn.RemoteAddr(), req.Host)
	return
}

type Http2Server struct {
	Base      *ProxyServer
	Handler   http.Handler
	TLSConfig *tls.Config
}

func NewHttp2Server(base *ProxyServer) *Http2Server {
	return &Http2Server{Base: base}
}

func (s *Http2Server) ListenAndServeTLS(config *tls.Config) error {
	srv := http.Server{
		Addr:      s.Base.Node.Addr,
		Handler:   s.Handler,
		TLSConfig: config,
	}
	if srv.Handler == nil {
		srv.Handler = http.HandlerFunc(s.HandleRequest)
	}
	http2.ConfigureServer(&srv, nil)
	return srv.ListenAndServeTLS("", "")
}

// Default HTTP2 server handler
func (s *Http2Server) HandleRequest(w http.ResponseWriter, req *http.Request) {
	target := req.Header.Get("Gost-Target")
	if target == "" {
		target = req.Host
	}
	glog.V(LINFO).Infof("[http2] %s %s - %s %s", req.Method, req.RemoteAddr, target, req.Proto)
	if glog.V(LDEBUG) {
		dump, _ := httputil.DumpRequest(req, false)
		glog.Infoln(string(dump))
	}

	w.Header().Set("Proxy-Agent", "gost/"+Version)

	// HTTP2 as transport
	if req.Header.Get("Proxy-Switch") == "gost" {
		conn, err := s.Upgrade(w, req)
		if err != nil {
			glog.V(LINFO).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
			return
		}
		glog.V(LINFO).Infof("[http2] %s - %s : switch to HTTP2 transport mode OK", req.RemoteAddr, target)
		s.Base.handleConn(conn)
		return
	}

	valid := false
	u, p, _ := basicProxyAuth(req.Header.Get("Proxy-Authorization"))
	for _, user := range s.Base.Node.Users {
		username := user.Username()
		password, _ := user.Password()
		if (u == username && p == password) ||
			(u == username && password == "") ||
			(username == "" && p == password) {
			valid = true
			break
		}
	}
	if len(s.Base.Node.Users) > 0 && !valid {
		glog.V(LWARNING).Infof("[http2] %s <- %s : proxy authentication required", req.RemoteAddr, target)
		w.WriteHeader(http.StatusProxyAuthRequired)
		return
	}

	req.Header.Del("Proxy-Authorization")

	c, err := s.Base.Chain.Dial(target)
	if err != nil {
		glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	defer c.Close()

	glog.V(LINFO).Infof("[http2] %s <-> %s", req.RemoteAddr, target)

	if req.Method == http.MethodConnect {
		w.WriteHeader(http.StatusOK)
		if fw, ok := w.(http.Flusher); ok {
			fw.Flush()
		}

		// compatible with HTTP1.x
		if hj, ok := w.(http.Hijacker); ok && req.ProtoMajor == 1 {
			// we take over the underly connection
			conn, _, err := hj.Hijack()
			if err != nil {
				glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			defer conn.Close()

			s.Base.transport(conn, c)
			return
		}

		errc := make(chan error, 2)

		go func() {
			_, err := io.Copy(c, req.Body)
			errc <- err
		}()
		go func() {
			_, err := io.Copy(flushWriter{w}, c)
			errc <- err
		}()

		select {
		case <-errc:
			// glog.V(LWARNING).Infoln("exit", err)
		}
		glog.V(LINFO).Infof("[http2] %s >-< %s", req.RemoteAddr, target)
		return
	}

	req.Header.Set("Connection", "Keep-Alive")
	if err = req.Write(c); err != nil {
		glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
		return
	}

	resp, err := http.ReadResponse(bufio.NewReader(c), req)
	if err != nil {
		glog.V(LWARNING).Infoln(err)
		return
	}
	defer resp.Body.Close()

	for k, v := range resp.Header {
		for _, vv := range v {
			w.Header().Add(k, vv)
		}
	}
	w.WriteHeader(resp.StatusCode)
	if _, err := io.Copy(flushWriter{w}, resp.Body); err != nil {
		glog.V(LWARNING).Infof("[http2] %s <- %s : %s", req.RemoteAddr, target, err)
	}

	glog.V(LINFO).Infof("[http2] %s >-< %s", req.RemoteAddr, target)
}

// Upgrade upgrade an HTTP2 request to a bidirectional connection that preparing for tunneling other protocol, just like a websocket connection.
func (s *Http2Server) Upgrade(w http.ResponseWriter, r *http.Request) (net.Conn, error) {
	if r.Method != http.MethodConnect {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return nil, errors.New("Method not allowed")
	}

	w.WriteHeader(http.StatusOK)

	if fw, ok := w.(http.Flusher); ok {
		fw.Flush()
	}

	conn := &http2Conn{r: r.Body, w: flushWriter{w}}
	conn.remoteAddr, _ = net.ResolveTCPAddr("tcp", r.RemoteAddr)
	conn.localAddr, _ = net.ResolveTCPAddr("tcp", r.Host)
	return conn, nil
}

// HTTP2 client connection, wrapped up just like a net.Conn
type http2Conn struct {
	r          io.Reader
	w          io.Writer
	remoteAddr net.Addr
	localAddr  net.Addr
}

func (c *http2Conn) Read(b []byte) (n int, err error) {
	return c.r.Read(b)
}

func (c *http2Conn) Write(b []byte) (n int, err error) {
	return c.w.Write(b)
}

func (c *http2Conn) Close() (err error) {
	if rc, ok := c.r.(io.Closer); ok {
		err = rc.Close()
	}
	if w, ok := c.w.(io.Closer); ok {
		err = w.Close()
	}
	return
}

func (c *http2Conn) LocalAddr() net.Addr {
	return c.localAddr
}

func (c *http2Conn) RemoteAddr() net.Addr {
	return c.remoteAddr
}

func (c *http2Conn) SetDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (c *http2Conn) SetReadDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (c *http2Conn) SetWriteDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

type flushWriter struct {
	w io.Writer
}

func (fw flushWriter) Write(p []byte) (n int, err error) {
	defer func() {
		if r := recover(); r != nil {
			if s, ok := r.(string); ok {
				err = errors.New(s)
				return
			}
			err = r.(error)
		}
	}()

	n, err = fw.w.Write(p)
	if err != nil {
		// glog.V(LWARNING).Infoln("flush writer:", err)
		return
	}
	if f, ok := fw.w.(http.Flusher); ok {
		f.Flush()
	}
	return
}
<commit_msg>add "Proxy-Authenticate" header for http2<commit_after>package gost

import (
	"bufio"
	"crypto/tls"
	"encoding/base64"
	"github.com/golang/glog"
	"golang.org/x/net/http2"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	//"strings"
	"errors"
	"time"
)

type HttpServer struct {
	conn net.Conn
	Base *ProxyServer
}

func NewHttpServer(conn net.Conn, base *ProxyServer) *HttpServer {
	return &HttpServer{
		conn: conn,
		Base: base,
	}
}

// Default HTTP server handler
func (s *HttpServer) HandleRequest(req *http.Request) {
	glog.V(LINFO).Infof("[http] %s %s - %s %s", req.Method, s.conn.RemoteAddr(), req.Host, req.Proto)

	if glog.V(LDEBUG) {
		dump, _ := httputil.DumpRequest(req, false)
		glog.Infoln(string(dump))
	}

	if req.Method == "PRI" && req.ProtoMajor == 2 {
		glog.V(LWARNING).Infof("[http] %s <- %s : Not an HTTP2 server", s.conn.RemoteAddr(), req.Host)
		resp := "HTTP/1.1 400 Bad Request\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n"
		s.conn.Write([]byte(resp))
		return
	}

	valid := false
	u, p, _ := basicProxyAuth(req.Header.Get("Proxy-Authorization"))
	glog.V(LINFO).Infoln(u, p)
	for _, user := range s.Base.Node.Users {
		username := user.Username()
		password, _ := user.Password()
		if (u == username && p == password) ||
			(u == username && password == "") ||
			(username == "" && p == password) {
			valid = true
			break
		}
	}

	if len(s.Base.Node.Users) > 0 && !valid {
		glog.V(LWARNING).Infof("[http] %s <- %s : proxy authentication required", s.conn.RemoteAddr(), req.Host)
		resp := "HTTP/1.1 407 Proxy Authentication Required\r\n" +
			"Proxy-Authenticate: Basic realm=\"gost\"\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n"
		s.conn.Write([]byte(resp))
		return
	}

	req.Header.Del("Proxy-Authorization")

	// forward http request
	lastNode := s.Base.Chain.lastNode
	if lastNode != nil && (lastNode.Protocol == "http" || lastNode.Protocol == "") {
		s.forwardRequest(req)
		return
	}

	c, err := s.Base.Chain.Dial(req.Host)
	if err != nil {
		glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), req.Host, err)

		b := []byte("HTTP/1.1 503 Service unavailable\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
		glog.V(LDEBUG).Infof("[http] %s <- %s\n%s", s.conn.RemoteAddr(), req.Host, string(b))
		s.conn.Write(b)
		return
	}
	defer c.Close()

	if req.Method == http.MethodConnect {
		b := []byte("HTTP/1.1 200 Connection established\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
		glog.V(LDEBUG).Infof("[http] %s <- %s\n%s", s.conn.RemoteAddr(), req.Host, string(b))
		s.conn.Write(b)
	} else {
		req.Header.Del("Proxy-Connection")
		req.Header.Set("Connection", "Keep-Alive")

		if err = req.Write(c); err != nil {
			glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), req.Host, err)
			return
		}
	}

	glog.V(LINFO).Infof("[http] %s <-> %s", s.conn.RemoteAddr(), req.Host)
	s.Base.transport(s.conn, c)
	glog.V(LINFO).Infof("[http] %s >-< %s", s.conn.RemoteAddr(), req.Host)
}

func (s *HttpServer) forwardRequest(req *http.Request) {
	last := s.Base.Chain.lastNode
	if last == nil {
		return
	}
	cc, err := s.Base.Chain.GetConn()
	if err != nil {
		glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), last.Addr, err)

		b := []byte("HTTP/1.1 503 Service unavailable\r\n" +
			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
		glog.V(LDEBUG).Infof("[http] %s <- %s\n%s", s.conn.RemoteAddr(), last.Addr, string(b))
		s.conn.Write(b)
		return
	}
	defer cc.Close()

	if len(last.Users) > 0 {
		user := last.Users[0]
		s := user.String()
		if _, set := user.Password(); !set {
			s += ":"
		}
		req.Header.Set("Proxy-Authorization",
			"Basic "+base64.StdEncoding.EncodeToString([]byte(s)))
	}

	cc.SetWriteDeadline(time.Now().Add(WriteTimeout))
	if err = req.WriteProxy(cc); err != nil {
		glog.V(LWARNING).Infof("[http] %s -> %s : %s", s.conn.RemoteAddr(), req.Host, err)
		return
	}
	cc.SetWriteDeadline(time.Time{})

	glog.V(LINFO).Infof("[http] %s <-> %s", s.conn.RemoteAddr(), req.Host)
	s.Base.transport(s.conn, cc)
	glog.V(LINFO).Infof("[http] %s >-< %s", s.conn.RemoteAddr(), req.Host)
	return
}

type Http2Server struct {
	Base      *ProxyServer
	Handler   http.Handler
	TLSConfig *tls.Config
}

func NewHttp2Server(base *ProxyServer) *Http2Server {
	return &Http2Server{Base: base}
}

func (s *Http2Server) ListenAndServeTLS(config *tls.Config) error {
	srv := http.Server{
		Addr:      s.Base.Node.Addr,
		Handler:   s.Handler,
		TLSConfig: config,
	}
	if srv.Handler == nil {
		srv.Handler = http.HandlerFunc(s.HandleRequest)
	}
	http2.ConfigureServer(&srv, nil)
	return srv.ListenAndServeTLS("", "")
}

// Default HTTP2 server handler
func (s *Http2Server) HandleRequest(w http.ResponseWriter, req *http.Request) {
	target := req.Header.Get("Gost-Target")
	if target == "" {
		target = req.Host
	}
	glog.V(LINFO).Infof("[http2] %s %s - %s %s", req.Method, req.RemoteAddr, target, req.Proto)
	if glog.V(LDEBUG) {
		dump, _ := httputil.DumpRequest(req, false)
		glog.Infoln(string(dump))
	}

	w.Header().Set("Proxy-Agent", "gost/"+Version)

	// HTTP2 as transport
	if req.Header.Get("Proxy-Switch") == "gost" {
		conn, err := s.Upgrade(w, req)
		if err != nil {
			glog.V(LINFO).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
			return
		}
		glog.V(LINFO).Infof("[http2] %s - %s : switch to HTTP2 transport mode OK", req.RemoteAddr, target)
		s.Base.handleConn(conn)
		return
	}

	valid := false
	u, p, _ := basicProxyAuth(req.Header.Get("Proxy-Authorization"))
	for _, user := range s.Base.Node.Users {
		username := user.Username()
		password, _ := user.Password()
		if (u == username && p == password) ||
			(u == username && password == "") ||
			(username == "" && p == password) {
			valid = true
			break
		}
	}
	if len(s.Base.Node.Users) > 0 && !valid {
		glog.V(LWARNING).Infof("[http2] %s <- %s : proxy authentication required", req.RemoteAddr, target)
		w.Header().Set("Proxy-Authenticate", "Basic realm=\"gost\"")
		w.WriteHeader(http.StatusProxyAuthRequired)
		return
	}

	req.Header.Del("Proxy-Authorization")

	c, err := s.Base.Chain.Dial(target)
	if err != nil {
		glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	defer c.Close()

	glog.V(LINFO).Infof("[http2] %s <-> %s", req.RemoteAddr, target)

	if req.Method == http.MethodConnect {
		w.WriteHeader(http.StatusOK)
		if fw, ok := w.(http.Flusher); ok {
			fw.Flush()
		}

		// compatible with HTTP1.x
		if hj, ok := w.(http.Hijacker); ok && req.ProtoMajor == 1 {
			// we take over the underly connection
			conn, _, err := hj.Hijack()
			if err != nil {
				glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			defer conn.Close()

			s.Base.transport(conn, c)
			return
		}

		errc := make(chan error, 2)

		go func() {
			_, err := io.Copy(c, req.Body)
			errc <- err
		}()
		go func() {
			_, err := io.Copy(flushWriter{w}, c)
			errc <- err
		}()

		select {
		case <-errc:
			// glog.V(LWARNING).Infoln("exit", err)
		}
		glog.V(LINFO).Infof("[http2] %s >-< %s", req.RemoteAddr, target)
		return
	}

	req.Header.Set("Connection", "Keep-Alive")
	if err = req.Write(c); err != nil {
		glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, target, err)
		return
	}

	resp, err := http.ReadResponse(bufio.NewReader(c), req)
	if err != nil {
		glog.V(LWARNING).Infoln(err)
		return
	}
	defer resp.Body.Close()

	for k, v := range resp.Header {
		for _, vv := range v {
			w.Header().Add(k, vv)
		}
	}
	w.WriteHeader(resp.StatusCode)
	if _, err := io.Copy(flushWriter{w}, resp.Body); err != nil {
		glog.V(LWARNING).Infof("[http2] %s <- %s : %s", req.RemoteAddr, target, err)
	}

	glog.V(LINFO).Infof("[http2] %s >-< %s", req.RemoteAddr, target)
}

// Upgrade upgrade an HTTP2 request to a bidirectional connection that preparing for tunneling other protocol, just like a websocket connection.
func (s *Http2Server) Upgrade(w http.ResponseWriter, r *http.Request) (net.Conn, error) {
	if r.Method != http.MethodConnect {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return nil, errors.New("Method not allowed")
	}

	w.WriteHeader(http.StatusOK)

	if fw, ok := w.(http.Flusher); ok {
		fw.Flush()
	}

	conn := &http2Conn{r: r.Body, w: flushWriter{w}}
	conn.remoteAddr, _ = net.ResolveTCPAddr("tcp", r.RemoteAddr)
	conn.localAddr, _ = net.ResolveTCPAddr("tcp", r.Host)
	return conn, nil
}

// HTTP2 client connection, wrapped up just like a net.Conn
type http2Conn struct {
	r          io.Reader
	w          io.Writer
	remoteAddr net.Addr
	localAddr  net.Addr
}

func (c *http2Conn) Read(b []byte) (n int, err error) {
	return c.r.Read(b)
}

func (c *http2Conn) Write(b []byte) (n int, err error) {
	return c.w.Write(b)
}

func (c *http2Conn) Close() (err error) {
	if rc, ok := c.r.(io.Closer); ok {
		err = rc.Close()
	}
	if w, ok := c.w.(io.Closer); ok {
		err = w.Close()
	}
	return
}

func (c *http2Conn) LocalAddr() net.Addr {
	return c.localAddr
}

func (c *http2Conn) RemoteAddr() net.Addr {
	return c.remoteAddr
}

func (c *http2Conn) SetDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (c *http2Conn) SetReadDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (c *http2Conn) SetWriteDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "http2", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

type flushWriter struct {
	w io.Writer
}

func (fw flushWriter) Write(p []byte) (n int, err error) {
	defer func() {
		if r := recover(); r != nil {
			if s, ok := r.(string); ok {
				err = errors.New(s)
				return
			}
			err = r.(error)
		}
	}()

	n, err = fw.w.Write(p)
	if err != nil {
		// glog.V(LWARNING).Infoln("flush writer:", err)
		return
	}
	if f, ok := fw.w.(http.Flusher); ok {
		f.Flush()
	}
	return
}
<|endoftext|>