<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"runtime/debug"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	_ "net/http/pprof"

	"github.com/xindong/frontd/aes256cbc"
)

const (
	// max open file should at least be
	_MaxOpenfile              = uint64(1024 * 1024 * 1024)
	_MaxBackendAddrCacheCount = 1024 * 1024
	_DefaultPort              = 4043
	_MTU                      = 1500
)

var (
	_hdrCipherOrigin   = []byte("x-cipher-origin")
	_hdrForwardedFor   = []byte("x-forwarded-for")
	_maxHTTPHeaderSize = 4096 * 2
	_minHTTPHeaderSize = 32
)

var (
	_SecretPassphase []byte
	_Aes256CBC       = aes256cbc.New()
)

var (
	_BackendAddrCacheMutex sync.Mutex
	_BackendAddrCache      atomic.Value
	_BufioReaderPool       sync.Pool
)

var (
	_BackendDialTimeout = 5
)

type backendAddrMap map[string][]byte

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	os.Setenv("GOTRACEBACK", "crash")

	_BackendAddrCache.Store(make(backendAddrMap))

	var lim syscall.Rlimit
	syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim)
	if lim.Cur < _MaxOpenfile || lim.Max < _MaxOpenfile {
		lim.Cur = _MaxOpenfile
		lim.Max = _MaxOpenfile
		syscall.Setrlimit(syscall.RLIMIT_NOFILE, &lim)
	}

	_SecretPassphase = []byte(os.Getenv("SECRET"))

	mhs, err := strconv.Atoi(os.Getenv("MAX_HTTP_HEADER_SIZE"))
	if err == nil && mhs > _minHTTPHeaderSize {
		_maxHTTPHeaderSize = mhs
	}

	bt, err := strconv.Atoi(os.Getenv("BACKEND_TIMEOUT"))
	if err == nil && bt > 0 {
		_BackendDialTimeout = bt
	}

	pprofPort, err := strconv.Atoi(os.Getenv("PPROF_PORT"))
	if err == nil && pprofPort > 0 && pprofPort <= 65535 {
		go func() {
			log.Println(http.ListenAndServe(":"+strconv.Itoa(pprofPort), nil))
		}()
	}

	listenAndServe()

	log.Println("Exiting")
}

func listenAndServe() {
	l, err := net.Listen("tcp", ":"+strconv.Itoa(_DefaultPort))
	if err != nil {
		log.Fatal(err)
	}
	defer l.Close()
	var tempDelay time.Duration
	for {
		conn, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				time.Sleep(tempDelay)
				continue
			}
			log.Fatal(err)
		}
		tempDelay = 0
		go handleConn(conn)
	}
}

func handleConn(c net.Conn) {
	defer func() {
		c.Close()
		if r := recover(); r != nil {
			log.Println("Recovered in", r, ":", string(debug.Stack()))
		}
	}()

	rdr, ok := _BufioReaderPool.Get().(*bufio.Reader)
	if ok {
		rdr.Reset(c)
	} else {
		rdr = bufio.NewReader(c)
	}
	defer _BufioReaderPool.Put(rdr)

	addr, err := handleBinaryHdr(rdr, c)
	if err != nil {
		log.Println("x", err)
		return
	}

	var header *bytes.Buffer
	if addr == nil {
		// Read first line
		line, isPrefix, err := rdr.ReadLine()
		if err != nil || isPrefix {
			log.Println(err)
			writeErrCode(c, []byte("4104"), false)
			return
		}

		cipherAddr := line

		// check if it's HTTP request
		if bytes.Contains(line, []byte("HTTP")) {
			header = bytes.NewBuffer(line)
			header.Write([]byte("\n"))

			cipherAddr, err = handleHTTPHdr(rdr, c, header)
			if err != nil {
				log.Println(err)
				return
			}
		}

		// base64 decode
		dbuf := make([]byte, base64.StdEncoding.DecodedLen(len(cipherAddr)))
		n, err := base64.StdEncoding.Decode(dbuf, cipherAddr)
		if err != nil {
			writeErrCode(c, []byte("4106"), false)
			return
		}

		addr, err = backendAddrDecrypt(dbuf[:n])
		if err != nil {
			writeErrCode(c, []byte("4106"), false)
			return
		}
	}

	// TODO: check if addr is allowed

	// Build tunnel
	err = tunneling(string(addr), rdr, c, header)
	if err != nil {
		log.Println(err)
	}
}

func writeErrCode(c net.Conn, errCode []byte, httpws bool) {
	switch httpws {
	case true:
		fmt.Fprintf(c, "HTTP/1.1 %s Error\nConnection: Close", errCode)
	default:
		c.Write(errCode)
	}
}

func handleBinaryHdr(rdr *bufio.Reader, c net.Conn) (addr []byte, err error) {
	// use binary protocol if first byte is 0x00
	b, err := rdr.ReadByte()
	if err != nil {
		// TODO: how to cause error to test this?
		writeErrCode(c, []byte("4103"), false)
		return nil, err
	}
	if b == byte(0x00) {
		// binary protocol
		blen, err := rdr.ReadByte()
		if err != nil || blen == 0 {
			writeErrCode(c, []byte("4103"), false)
			return nil, err
		}
		p := make([]byte, blen)
		n, err := io.ReadFull(rdr, p)
		if n != int(blen) {
			// TODO: how to cause error to test this?
			writeErrCode(c, []byte("4109"), false)
			return nil, err
		}

		// decrypt
		addr, err := backendAddrDecrypt(p)
		if err != nil {
			writeErrCode(c, []byte("4106"), false)
			return nil, err
		}

		return addr, err
	}

	rdr.UnreadByte()
	return nil, nil
}

func handleHTTPHdr(rdr *bufio.Reader, c net.Conn, header *bytes.Buffer) (addr []byte, err error) {
	hdrXff := "X-Forwarded-For: " + ipAddrFromRemoteAddr(c.RemoteAddr().String())

	var cipherAddr []byte
	for {
		line, isPrefix, err := rdr.ReadLine()
		if err != nil || isPrefix {
			log.Println(err)
			writeErrCode(c, []byte("4107"), true)
			return nil, err
		}

		if bytes.HasPrefix(bytes.ToLower(line), _hdrCipherOrigin) {
			// copy instead of point
			cipherAddr = []byte(string(bytes.TrimSpace(line[(len(_hdrCipherOrigin) + 1):])))
			continue
		}

		if bytes.HasPrefix(bytes.ToLower(line), _hdrForwardedFor) {
			hdrXff = hdrXff + ", " + string(bytes.TrimSpace(line[(len(_hdrForwardedFor)+1):]))
			continue
		}

		if len(bytes.TrimSpace(line)) == 0 {
			// end of HTTP header
			if len(cipherAddr) == 0 {
				writeErrCode(c, []byte("4108"), true)
				return nil, errors.New("empty http cipher address header")
			}
			if len(hdrXff) > 0 {
				header.Write([]byte(hdrXff))
				header.Write([]byte("\n"))
			}
			header.Write(line)
			header.Write([]byte("\n"))
			break
		}

		header.Write(line)
		header.Write([]byte("\n"))

		if header.Len() > _maxHTTPHeaderSize {
			writeErrCode(c, []byte("4108"), true)
			return nil, errors.New("http header size overflowed")
		}
	}

	return cipherAddr, nil
}

// tunneling to backend
func tunneling(addr string, rdr *bufio.Reader, c net.Conn, header *bytes.Buffer) error {
	backend, err := dialTimeout("tcp", addr, time.Second*time.Duration(_BackendDialTimeout))
	if err != nil {
		// handle error
		switch err := err.(type) {
		case net.Error:
			if err.Timeout() {
				writeErrCode(c, []byte("4101"), false)
				return err
			}
		}
		writeErrCode(c, []byte("4102"), false)
		return err
	}
	defer backend.Close()

	if header != nil {
		header.WriteTo(backend)
	}

	// Start transfering data
	ch := make(chan struct{}, 2)
	go pipe(c, backend, ch)
	go pipe(backend, rdr, ch)
	<-ch
	return nil
}

func dialTimeout(network, address string, timeout time.Duration) (conn net.Conn, err error) {
	m := int(timeout / time.Second)
	for i := 0; i < m; i++ {
		conn, err = net.DialTimeout(network, address, timeout)
		if err == nil || !strings.Contains(err.Error(), "can't assign requested address") {
			break
		}
		time.Sleep(time.Second)
	}
	return
}

func backendAddrDecrypt(key []byte) ([]byte, error) {
	// Try to check cache
	m1 := _BackendAddrCache.Load().(backendAddrMap)
	k1 := string(key)
	addr, ok := m1[k1]
	if ok {
		return addr, nil
	}

	// Try to decrypt it (AES)
	addr, err := _Aes256CBC.Decrypt(_SecretPassphase, key)
	if err != nil {
		return nil, err
	}

	backendAddrList(k1, addr)
	return addr, nil
}

func backendAddrList(key string, val []byte) {
	_BackendAddrCacheMutex.Lock()
	defer _BackendAddrCacheMutex.Unlock()

	m1 := _BackendAddrCache.Load().(backendAddrMap)
	// double check
	if _, ok := m1[key]; ok {
		return
	}

	m2 := make(backendAddrMap)
	// flush cache if there is way too many
	if len(m1) < _MaxBackendAddrCacheCount {
		// copy-on-write
		for k, v := range m1 {
			m2[k] = v // copy all data from the current object to the new one
		}
	}
	m2[key] = val
	_BackendAddrCache.Store(m2) // atomically replace the current object with the new one
}

// Request.RemoteAddress contains port, which we want to remove i.e.:
// "[::1]:58292" => "[::1]"
func ipAddrFromRemoteAddr(s string) string {
	idx := strings.LastIndex(s, ":")
	if idx == -1 {
		return s
	}
	return s[:idx]
}

// pipe upstream and downstream
func pipe(dst io.Writer, src io.Reader, ch chan struct{}) {
	defer func() {
		if r := recover(); r != nil {
			log.Println("Recovered in", r, ":", string(debug.Stack()))
		}
	}()
	defer func() {
		ch <- struct{}{}
	}()

	_, err := io.Copy(dst, src)

	switch err {
	case io.EOF:
		err = nil
		return
	case nil:
		return
	}
	// log.Println("pipe:", n, err)
}
<commit_msg>add connection read timeout. improving pipe closed<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"runtime/debug"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	_ "net/http/pprof"

	"github.com/xindong/frontd/aes256cbc"
)

const (
	// max open file should at least be
	_MaxOpenfile              = uint64(1024 * 1024 * 1024)
	_MaxBackendAddrCacheCount = 1024 * 1024
	_DefaultPort              = 4043
	_MTU                      = 1500
)

var (
	_hdrCipherOrigin   = []byte("x-cipher-origin")
	_hdrForwardedFor   = []byte("x-forwarded-for")
	_maxHTTPHeaderSize = 4096 * 2
	_minHTTPHeaderSize = 32
)

var (
	_SecretPassphase []byte
	_Aes256CBC       = aes256cbc.New()
)

var (
	_BackendAddrCacheMutex sync.Mutex
	_BackendAddrCache      atomic.Value
	_BufioReaderPool       sync.Pool
)

var (
	_BackendDialTimeout = 5
	_ConnReadTimeout    = time.Duration(300)
)

type backendAddrMap map[string][]byte

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	os.Setenv("GOTRACEBACK", "crash")

	_BackendAddrCache.Store(make(backendAddrMap))

	var lim syscall.Rlimit
	syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim)
	if lim.Cur < _MaxOpenfile || lim.Max < _MaxOpenfile {
		lim.Cur = _MaxOpenfile
		lim.Max = _MaxOpenfile
		syscall.Setrlimit(syscall.RLIMIT_NOFILE, &lim)
	}

	_SecretPassphase = []byte(os.Getenv("SECRET"))

	mhs, err := strconv.Atoi(os.Getenv("MAX_HTTP_HEADER_SIZE"))
	if err == nil && mhs > _minHTTPHeaderSize {
		_maxHTTPHeaderSize = mhs
	}

	bt, err := strconv.Atoi(os.Getenv("BACKEND_TIMEOUT"))
	if err == nil && bt > 0 {
		_BackendDialTimeout = bt
	}

	connReadTimeout, err := strconv.Atoi(os.Getenv("CONN_READ_TIMEOUT"))
	if err == nil && connReadTimeout >= 0 {
		_ConnReadTimeout = time.Duration(connReadTimeout)
	}

	pprofPort, err := strconv.Atoi(os.Getenv("PPROF_PORT"))
	if err == nil && pprofPort > 0 && pprofPort <= 65535 {
		go func() {
			log.Println(http.ListenAndServe(":"+strconv.Itoa(pprofPort), nil))
		}()
	}

	listenAndServe()

	log.Println("Exiting")
}

func listenAndServe() {
	l, err := net.Listen("tcp", ":"+strconv.Itoa(_DefaultPort))
	if err != nil {
		log.Fatal(err)
	}
	defer l.Close()
	var tempDelay time.Duration
	for {
		conn, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				time.Sleep(tempDelay)
				continue
			}
			log.Fatal(err)
		}
		tempDelay = 0
		go handleConn(conn)
	}
}

func handleConn(c net.Conn) {
	defer func() {
		c.Close()
		if r := recover(); r != nil {
			log.Println("Recovered in", r, ":", string(debug.Stack()))
		}
	}()

	rdr, ok := _BufioReaderPool.Get().(*bufio.Reader)
	if ok {
		rdr.Reset(c)
	} else {
		rdr = bufio.NewReader(c)
	}
	defer _BufioReaderPool.Put(rdr)

	addr, err := handleBinaryHdr(rdr, c)
	if err != nil {
		log.Println("x", err)
		return
	}

	var header *bytes.Buffer
	if addr == nil {
		// Read first line
		line, isPrefix, err := rdr.ReadLine()
		if err != nil || isPrefix {
			log.Println(err)
			writeErrCode(c, []byte("4104"), false)
			return
		}

		cipherAddr := line

		// check if it's HTTP request
		if bytes.Contains(line, []byte("HTTP")) {
			header = bytes.NewBuffer(line)
			header.Write([]byte("\n"))

			cipherAddr, err = handleHTTPHdr(rdr, c, header)
			if err != nil {
				log.Println(err)
				return
			}
		}

		// base64 decode
		dbuf := make([]byte, base64.StdEncoding.DecodedLen(len(cipherAddr)))
		n, err := base64.StdEncoding.Decode(dbuf, cipherAddr)
		if err != nil {
			writeErrCode(c, []byte("4106"), false)
			return
		}

		addr, err = backendAddrDecrypt(dbuf[:n])
		if err != nil {
			writeErrCode(c, []byte("4106"), false)
			return
		}
	}

	// TODO: check if addr is allowed

	// Build tunnel
	err = tunneling(string(addr), rdr, c, header)
	if err != nil {
		log.Println(err)
	}
}

func writeErrCode(c net.Conn, errCode []byte, httpws bool) {
	switch httpws {
	case true:
		fmt.Fprintf(c, "HTTP/1.1 %s Error\nConnection: Close", errCode)
	default:
		c.Write(errCode)
	}
}

func handleBinaryHdr(rdr *bufio.Reader, c net.Conn) (addr []byte, err error) {
	// use binary protocol if first byte is 0x00
	b, err := rdr.ReadByte()
	if err != nil {
		// TODO: how to cause error to test this?
		writeErrCode(c, []byte("4103"), false)
		return nil, err
	}
	if b == byte(0x00) {
		// binary protocol
		blen, err := rdr.ReadByte()
		if err != nil || blen == 0 {
			writeErrCode(c, []byte("4103"), false)
			return nil, err
		}
		p := make([]byte, blen)
		n, err := io.ReadFull(rdr, p)
		if n != int(blen) {
			// TODO: how to cause error to test this?
			writeErrCode(c, []byte("4109"), false)
			return nil, err
		}

		// decrypt
		addr, err := backendAddrDecrypt(p)
		if err != nil {
			writeErrCode(c, []byte("4106"), false)
			return nil, err
		}

		return addr, err
	}

	rdr.UnreadByte()
	return nil, nil
}

func handleHTTPHdr(rdr *bufio.Reader, c net.Conn, header *bytes.Buffer) (addr []byte, err error) {
	hdrXff := "X-Forwarded-For: " + ipAddrFromRemoteAddr(c.RemoteAddr().String())

	var cipherAddr []byte
	for {
		line, isPrefix, err := rdr.ReadLine()
		if err != nil || isPrefix {
			log.Println(err)
			writeErrCode(c, []byte("4107"), true)
			return nil, err
		}

		if bytes.HasPrefix(bytes.ToLower(line), _hdrCipherOrigin) {
			// copy instead of point
			cipherAddr = []byte(string(bytes.TrimSpace(line[(len(_hdrCipherOrigin) + 1):])))
			continue
		}

		if bytes.HasPrefix(bytes.ToLower(line), _hdrForwardedFor) {
			hdrXff = hdrXff + ", " + string(bytes.TrimSpace(line[(len(_hdrForwardedFor)+1):]))
			continue
		}

		if len(bytes.TrimSpace(line)) == 0 {
			// end of HTTP header
			if len(cipherAddr) == 0 {
				writeErrCode(c, []byte("4108"), true)
				return nil, errors.New("empty http cipher address header")
			}
			if len(hdrXff) > 0 {
				header.Write([]byte(hdrXff))
				header.Write([]byte("\n"))
			}
			header.Write(line)
			header.Write([]byte("\n"))
			break
		}

		header.Write(line)
		header.Write([]byte("\n"))

		if header.Len() > _maxHTTPHeaderSize {
			writeErrCode(c, []byte("4108"), true)
			return nil, errors.New("http header size overflowed")
		}
	}

	return cipherAddr, nil
}

// tunneling to backend
func tunneling(addr string, rdr *bufio.Reader, c net.Conn, header *bytes.Buffer) error {
	backend, err := dialTimeout("tcp", addr, time.Second*time.Duration(_BackendDialTimeout))
	if err != nil {
		// handle error
		switch err := err.(type) {
		case net.Error:
			if err.Timeout() {
				writeErrCode(c, []byte("4101"), false)
				return err
			}
		}
		writeErrCode(c, []byte("4102"), false)
		return err
	}
	defer backend.Close()

	if header != nil {
		header.WriteTo(backend)
	}

	// Start transfering data
	go pipe(c, backend, c, backend)
	pipe(backend, rdr, backend, c)

	return nil
}

func dialTimeout(network, address string, timeout time.Duration) (conn net.Conn, err error) {
	m := int(timeout / time.Second)
	for i := 0; i < m; i++ {
		conn, err = net.DialTimeout(network, address, timeout)
		if err == nil || !strings.Contains(err.Error(), "can't assign requested address") {
			break
		}
		time.Sleep(time.Second)
	}
	return
}

func backendAddrDecrypt(key []byte) ([]byte, error) {
	// Try to check cache
	m1 := _BackendAddrCache.Load().(backendAddrMap)
	k1 := string(key)
	addr, ok := m1[k1]
	if ok {
		return addr, nil
	}

	// Try to decrypt it (AES)
	addr, err := _Aes256CBC.Decrypt(_SecretPassphase, key)
	if err != nil {
		return nil, err
	}

	backendAddrList(k1, addr)
	return addr, nil
}

func backendAddrList(key string, val []byte) {
	_BackendAddrCacheMutex.Lock()
	defer _BackendAddrCacheMutex.Unlock()

	m1 := _BackendAddrCache.Load().(backendAddrMap)
	// double check
	if _, ok := m1[key]; ok {
		return
	}

	m2 := make(backendAddrMap)
	// flush cache if there is way too many
	if len(m1) < _MaxBackendAddrCacheCount {
		// copy-on-write
		for k, v := range m1 {
			m2[k] = v // copy all data from the current object to the new one
		}
	}
	m2[key] = val
	_BackendAddrCache.Store(m2) // atomically replace the current object with the new one
}

// Request.RemoteAddress contains port, which we want to remove i.e.:
// "[::1]:58292" => "[::1]"
func ipAddrFromRemoteAddr(s string) string {
	idx := strings.LastIndex(s, ":")
	if idx == -1 {
		return s
	}
	return s[:idx]
}

// pipe upstream and downstream
func pipe(dst io.Writer, src io.Reader, dstconn, srcconn net.Conn) {
	defer func() {
		if r := recover(); r != nil {
			log.Println("Recovered in", r, ":", string(debug.Stack()))
		}
	}()

	// only close dst when done
	defer dstconn.Close()

	buf := make([]byte, 32*1024)
	for {
		srcconn.SetReadDeadline(time.Now().Add(time.Second * _ConnReadTimeout))
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if ew != nil {
				break
			}
			if nr != nw {
				break
			}
		}
		if er == io.EOF {
			break
		}
		if er != nil {
			break
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"gopkg.in/tylerb/graceful.v1"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

const githubStatusContext = "review"

type IssueComment struct {
	IssueNumber   int
	Comment       string
	IsPullRequest bool
	Repository    Repository
}

type PullRequestEvent struct {
	IssueNumber int
	Action      string
	Repository  Repository
}

type Repository struct {
	Owner string
	Name  string
	URL   string
}

func main() {
	conf := NewConfig()
	githubClient := initGithubClient(conf.AccessToken)
	reposDir, err := ioutil.TempDir("", "github-review-helper")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(reposDir)

	git := NewGit(reposDir)

	mux := http.NewServeMux()
	mux.Handle("/", Handler(func(w http.ResponseWriter, r *http.Request) Response {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return ErrorResponse{err, http.StatusInternalServerError, "Failed to read the request's body"}
		}
		signature := r.Header.Get("X-Hub-Signature")
		hasSecret, err := hasSecret(body, signature, conf.Secret)
		if err != nil {
			return ErrorResponse{err, http.StatusInternalServerError, "Failed to check the signature"}
		}
		if !hasSecret {
			return ErrorResponse{nil, http.StatusBadRequest, "Bad X-Hub-Signature"}
		}
		eventType := r.Header.Get("X-Github-Event")
		switch eventType {
		case "issue_comment":
			return handleIssueComment(w, body, git, githubClient)
		case "pull_request":
			return handlePullRequest(w, body, git, githubClient)
		}
		return SuccessResponse{"Not an event I understand. Ignoring."}
	}))

	graceful.Run(fmt.Sprintf(":%d", conf.Port), 10*time.Second, mux)
}

func handleIssueComment(w http.ResponseWriter, body []byte, git Git, githubClient *github.Client) Response {
	issueComment, err := parseIssueComment(body)
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to parse the request's body"}
	}
	if !issueComment.IsPullRequest {
		return SuccessResponse{"Not a PR. Ignoring."}
	}
	switch issueComment.Comment {
	case "!squash":
		return handleSquash(w, issueComment, git, githubClient)
	}
	return SuccessResponse{"Not a command I understand. Ignoring."}
}

func handleSquash(w http.ResponseWriter, issueComment IssueComment, git Git, githubClient *github.Client) Response {
	pr, _, err := githubClient.PullRequests.Get(issueComment.Repository.Owner, issueComment.Repository.Name, issueComment.IssueNumber)
	if err != nil {
		message := fmt.Sprintf("Getting PR %s/%s#%d failed", issueComment.Repository.Owner, issueComment.Repository.Name, issueComment.IssueNumber)
		return ErrorResponse{err, http.StatusBadGateway, message}
	}
	log.Printf("Squashing %s that's going to be merged into %s\n", *pr.Head.Ref, *pr.Base.Ref)
	repo, err := git.GetUpdatedRepo(issueComment.Repository.URL, issueComment.Repository.Owner, issueComment.Repository.Name)
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to update the local repo"}
	}
	if err = repo.RebaseAutosquash(*pr.Base.SHA, *pr.Head.SHA); err != nil {
		log.Printf("Failed to autosquash the commits with an interactive rebase: %s. Setting a failure status.\n", err)
		_, _, err = githubClient.Repositories.CreateStatus(issueComment.Repository.Owner, issueComment.Repository.Name, *pr.Head.SHA, &github.RepoStatus{
			State:       github.String("failure"),
			Description: github.String("Failed to automatically squash the fixup! and squash! commits. Please squash manually"),
			Context:     github.String(githubStatusContext),
		})
		if err != nil {
			message := fmt.Sprintf("Failed to create a failure status for commit %s", *pr.Head.SHA)
			return ErrorResponse{err, http.StatusBadGateway, message}
		}
		return SuccessResponse{"Failed to autosquash the commits with an interactive rebase. Reported the failure."}
	}
	if err = repo.ForcePushHeadTo(*pr.Head.Ref); err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to push the squashed version"}
	}
	headSHA, err := repo.GetHeadSHA()
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to get the squashed branch's HEAD's SHA"}
	}
	_, _, err = githubClient.Repositories.CreateStatus(issueComment.Repository.Owner, issueComment.Repository.Name, headSHA, &github.RepoStatus{
		State:       github.String("success"),
		Description: github.String("All fixup! and squash! commits successfully squashed"),
		Context:     github.String(githubStatusContext),
	})
	if err != nil {
		message := fmt.Sprintf("Failed to create a success status for commit %s", headSHA)
		return ErrorResponse{err, http.StatusBadGateway, message}
	}
	return SuccessResponse{}
}

func handlePullRequest(w http.ResponseWriter, body []byte, git Git, githubClient *github.Client) Response {
	pullRequestEvent, err := parsePullRequestEvent(body)
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to parse the request's body"}
	}
	if !(pullRequestEvent.Action == "opened" || pullRequestEvent.Action == "synchronize") {
		return SuccessResponse{"PR not opened or synchronized. Ignoring."}
	}
	log.Printf("Checking for fixup commits for PR %s/%s#%d.\n", pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
	commits, _, err := githubClient.PullRequests.ListCommits(pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber, nil)
	if err != nil {
		message := fmt.Sprintf("Getting commits for PR %s/%s#%d failed", pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
		return ErrorResponse{err, http.StatusBadGateway, message}
	}
	if includesFixupCommits(commits) {
		pr, _, err := githubClient.PullRequests.Get(pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
		if err != nil {
			message := fmt.Sprintf("Getting PR %s/%s#%d failed", pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
			return ErrorResponse{err, http.StatusBadGateway, message}
		}
		_, _, err = githubClient.Repositories.CreateStatus(pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, *pr.Head.SHA, &github.RepoStatus{
			State:       github.String("pending"),
			Description: github.String("This PR needs to be squashed with !squash before merging"),
			Context:     github.String(githubStatusContext),
		})
		if err != nil {
			message := fmt.Sprintf("Failed to create a pending status for commit %s", *pr.Head.SHA)
			return ErrorResponse{err, http.StatusBadGateway, message}
		}
	}
	return SuccessResponse{}
}

func includesFixupCommits(commits []github.RepositoryCommit) bool {
	for _, commit := range commits {
		if strings.HasPrefix(*commit.Commit.Message, "fixup! ") || strings.HasPrefix(*commit.Commit.Message, "squash! ") {
			return true
		}
	}
	return false
}

func initGithubClient(accessToken string) *github.Client {
	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: accessToken},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	return github.NewClient(tc)
}

func parseIssueComment(body []byte) (IssueComment, error) {
	var message struct {
		Issue struct {
			Number      int `json:"Number"`
			PullRequest struct {
				URL string `json:"url"`
			} `json:"pull_request"`
		} `json:"issue"`
		Repository struct {
			Name  string `json:"name"`
			Owner struct {
				Login string `json:"login"`
			} `json:"owner"`
			SSHURL string `json:"ssh_url"`
		} `json:"repository"`
		Comment struct {
			Body string `json:"body"`
		} `json:"comment"`
	}
	err := json.Unmarshal(body, &message)
	if err != nil {
		return IssueComment{}, err
	}
	return IssueComment{
		IssueNumber:   message.Issue.Number,
		Comment:       message.Comment.Body,
		IsPullRequest: message.Issue.PullRequest.URL != "",
		Repository: Repository{
			Owner: message.Repository.Owner.Login,
			Name:  message.Repository.Name,
			URL:   message.Repository.SSHURL,
		},
	}, nil
}

func parsePullRequestEvent(body []byte) (PullRequestEvent, error) {
	var message struct {
		Action      string `json:"action"`
		Number      int    `json:"number"`
		PullRequest struct {
			URL string `json:"url"`
		} `json:"pull_request"`
		Repository struct {
			Name  string `json:"name"`
			Owner struct {
				Login string `json:"login"`
			} `json:"owner"`
			SSHURL string `json:"ssh_url"`
		} `json:"repository"`
	}
	err := json.Unmarshal(body, &message)
	if err != nil {
		return PullRequestEvent{}, err
	}
	return PullRequestEvent{
		IssueNumber: message.Number,
		Action:      message.Action,
		Repository: Repository{
			Owner: message.Repository.Owner.Login,
			Name:  message.Repository.Name,
			URL:   message.Repository.SSHURL,
		},
	}, nil
}

func hasSecret(message []byte, signature, key string) (bool, error) {
	var messageMACString string
	fmt.Sscanf(signature, "sha1=%s", &messageMACString)
	messageMAC, err := hex.DecodeString(messageMACString)
	if err != nil {
		return false, err
	}

	mac := hmac.New(sha1.New, []byte(key))
	mac.Write(message)
	expectedMAC := mac.Sum(nil)
	return hmac.Equal(messageMAC, expectedMAC), nil
}
<commit_msg>Make the squash status context more specific<commit_after>package main

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"gopkg.in/tylerb/graceful.v1"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

const githubStatusSquashContext = "review/squash"

type IssueComment struct {
	IssueNumber   int
	Comment       string
	IsPullRequest bool
	Repository    Repository
}

type PullRequestEvent struct {
	IssueNumber int
	Action      string
	Repository  Repository
}

type Repository struct {
	Owner string
	Name  string
	URL   string
}

func main() {
	conf := NewConfig()
	githubClient := initGithubClient(conf.AccessToken)
	reposDir, err := ioutil.TempDir("", "github-review-helper")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(reposDir)

	git := NewGit(reposDir)

	mux := http.NewServeMux()
	mux.Handle("/", Handler(func(w http.ResponseWriter, r *http.Request) Response {
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return ErrorResponse{err, http.StatusInternalServerError, "Failed to read the request's body"}
		}
		signature := r.Header.Get("X-Hub-Signature")
		hasSecret, err := hasSecret(body, signature, conf.Secret)
		if err != nil {
			return ErrorResponse{err, http.StatusInternalServerError, "Failed to check the signature"}
		}
		if !hasSecret {
			return ErrorResponse{nil, http.StatusBadRequest, "Bad X-Hub-Signature"}
		}
		eventType := r.Header.Get("X-Github-Event")
		switch eventType {
		case "issue_comment":
			return handleIssueComment(w, body, git, githubClient)
		case "pull_request":
			return handlePullRequest(w, body, git, githubClient)
		}
		return SuccessResponse{"Not an event I understand. Ignoring."}
	}))

	graceful.Run(fmt.Sprintf(":%d", conf.Port), 10*time.Second, mux)
}

func handleIssueComment(w http.ResponseWriter, body []byte, git Git, githubClient *github.Client) Response {
	issueComment, err := parseIssueComment(body)
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to parse the request's body"}
	}
	if !issueComment.IsPullRequest {
		return SuccessResponse{"Not a PR. Ignoring."}
	}
	switch issueComment.Comment {
	case "!squash":
		return handleSquash(w, issueComment, git, githubClient)
	}
	return SuccessResponse{"Not a command I understand. Ignoring."}
}

func handleSquash(w http.ResponseWriter, issueComment IssueComment, git Git, githubClient *github.Client) Response {
	pr, _, err := githubClient.PullRequests.Get(issueComment.Repository.Owner, issueComment.Repository.Name, issueComment.IssueNumber)
	if err != nil {
		message := fmt.Sprintf("Getting PR %s/%s#%d failed", issueComment.Repository.Owner, issueComment.Repository.Name, issueComment.IssueNumber)
		return ErrorResponse{err, http.StatusBadGateway, message}
	}
	log.Printf("Squashing %s that's going to be merged into %s\n", *pr.Head.Ref, *pr.Base.Ref)
	repo, err := git.GetUpdatedRepo(issueComment.Repository.URL, issueComment.Repository.Owner, issueComment.Repository.Name)
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to update the local repo"}
	}
	if err = repo.RebaseAutosquash(*pr.Base.SHA, *pr.Head.SHA); err != nil {
		log.Printf("Failed to autosquash the commits with an interactive rebase: %s. Setting a failure status.\n", err)
		_, _, err = githubClient.Repositories.CreateStatus(issueComment.Repository.Owner, issueComment.Repository.Name, *pr.Head.SHA, &github.RepoStatus{
			State:       github.String("failure"),
			Description: github.String("Failed to automatically squash the fixup! and squash! commits. Please squash manually"),
			Context:     github.String(githubStatusSquashContext),
		})
		if err != nil {
			message := fmt.Sprintf("Failed to create a failure status for commit %s", *pr.Head.SHA)
			return ErrorResponse{err, http.StatusBadGateway, message}
		}
		return SuccessResponse{"Failed to autosquash the commits with an interactive rebase. Reported the failure."}
	}
	if err = repo.ForcePushHeadTo(*pr.Head.Ref); err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to push the squashed version"}
	}
	headSHA, err := repo.GetHeadSHA()
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to get the squashed branch's HEAD's SHA"}
	}
	_, _, err = githubClient.Repositories.CreateStatus(issueComment.Repository.Owner, issueComment.Repository.Name, headSHA, &github.RepoStatus{
		State:       github.String("success"),
		Description: github.String("All fixup! and squash! commits successfully squashed"),
		Context:     github.String(githubStatusSquashContext),
	})
	if err != nil {
		message := fmt.Sprintf("Failed to create a success status for commit %s", headSHA)
		return ErrorResponse{err, http.StatusBadGateway, message}
	}
	return SuccessResponse{}
}

func handlePullRequest(w http.ResponseWriter, body []byte, git Git, githubClient *github.Client) Response {
	pullRequestEvent, err := parsePullRequestEvent(body)
	if err != nil {
		return ErrorResponse{err, http.StatusInternalServerError, "Failed to parse the request's body"}
	}
	if !(pullRequestEvent.Action == "opened" || pullRequestEvent.Action == "synchronize") {
		return SuccessResponse{"PR not opened or synchronized. Ignoring."}
	}
	log.Printf("Checking for fixup commits for PR %s/%s#%d.\n", pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
	commits, _, err := githubClient.PullRequests.ListCommits(pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber, nil)
	if err != nil {
		message := fmt.Sprintf("Getting commits for PR %s/%s#%d failed", pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
		return ErrorResponse{err, http.StatusBadGateway, message}
	}
	if includesFixupCommits(commits) {
		pr, _, err := githubClient.PullRequests.Get(pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
		if err != nil {
			message := fmt.Sprintf("Getting PR %s/%s#%d failed", pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, pullRequestEvent.IssueNumber)
			return ErrorResponse{err, http.StatusBadGateway, message}
		}
		_, _, err = githubClient.Repositories.CreateStatus(pullRequestEvent.Repository.Owner, pullRequestEvent.Repository.Name, *pr.Head.SHA, &github.RepoStatus{
			State:       github.String("pending"),
			Description: github.String("This PR needs to be squashed with !squash before merging"),
			Context:     github.String(githubStatusSquashContext),
		})
		if err != nil {
			message := fmt.Sprintf("Failed to create a pending status for commit %s", *pr.Head.SHA)
			return ErrorResponse{err, http.StatusBadGateway, message}
		}
	}
	return SuccessResponse{}
}

func includesFixupCommits(commits []github.RepositoryCommit) bool {
	for _, commit := range commits {
		if strings.HasPrefix(*commit.Commit.Message, "fixup! ") || strings.HasPrefix(*commit.Commit.Message, "squash! ") {
			return true
		}
	}
	return false
}

func initGithubClient(accessToken string) *github.Client {
	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: accessToken},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)
	return github.NewClient(tc)
}

func parseIssueComment(body []byte) (IssueComment, error) {
	var message struct {
		Issue struct {
			Number      int `json:"Number"`
			PullRequest struct {
				URL string `json:"url"`
			} `json:"pull_request"`
		} `json:"issue"`
		Repository struct {
			Name  string `json:"name"`
			Owner struct {
				Login string `json:"login"`
			} `json:"owner"`
			SSHURL string `json:"ssh_url"`
		} `json:"repository"`
		Comment struct {
			Body string `json:"body"`
		} `json:"comment"`
	}
	err := json.Unmarshal(body, &message)
	if err != nil {
		return IssueComment{}, err
	}
	return IssueComment{
		IssueNumber:   message.Issue.Number,
		Comment:       message.Comment.Body,
		IsPullRequest: message.Issue.PullRequest.URL != "",
		Repository: Repository{
			Owner: message.Repository.Owner.Login,
			Name:  message.Repository.Name,
			URL:   message.Repository.SSHURL,
		},
	}, nil
}

func parsePullRequestEvent(body []byte) (PullRequestEvent, error) {
	var message struct {
		Action      string `json:"action"`
		Number      int    `json:"number"`
		PullRequest struct {
			URL string `json:"url"`
		} `json:"pull_request"`
		Repository struct {
			Name  string `json:"name"`
			Owner struct {
				Login string `json:"login"`
			} `json:"owner"`
			SSHURL string `json:"ssh_url"`
		} `json:"repository"`
	}
	err := json.Unmarshal(body, &message)
	if err != nil {
		return PullRequestEvent{}, err
	}
	return PullRequestEvent{
		IssueNumber: message.Number,
		Action:      message.Action,
		Repository: Repository{
			Owner: message.Repository.Owner.Login,
			Name:  message.Repository.Name,
			URL:   message.Repository.SSHURL,
		},
	}, nil
}

func hasSecret(message []byte, signature, key string) (bool, error) {
	var messageMACString string
	fmt.Sscanf(signature, "sha1=%s", &messageMACString)
	messageMAC, err := hex.DecodeString(messageMACString)
	if err != nil {
		return false, err
	}

	mac := hmac.New(sha1.New, []byte(key))
	mac.Write(message)
	expectedMAC := mac.Sum(nil)
	return hmac.Equal(messageMAC, expectedMAC), nil
}
<|endoftext|>


<commit_before>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

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

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

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

*/
package main

import (
	"flag"
	"fmt"
	stdLog "log"
	"net/http"
	"os"
	"strings"

	_ "github.com/admpub/nging/application/library/sqlite"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/engine"
	"github.com/webx-top/echo/engine/standard"
	"github.com/webx-top/echo/handler/mvc/events"
	"github.com/webx-top/echo/middleware"
	"github.com/webx-top/echo/middleware/language"
	"github.com/webx-top/echo/middleware/render"
	"github.com/webx-top/echo/middleware/render/driver"
	"github.com/webx-top/echo/middleware/session"

	"github.com/admpub/log"
	"github.com/admpub/nging/application"
	"github.com/admpub/nging/application/library/config"
	"github.com/admpub/nging/application/library/cron"
	"github.com/admpub/nging/application/library/service"
)

var (
	//Version 版本号
	Version = `1.2.0`

	binData    bool
	staticMW   interface{}
	tmplMgr    driver.Manager
	langFSFunc func(dir string) http.FileSystem
)

func main() {
	config.DefaultCLIConfig.InitFlag()
	flag.Parse()

	if binData {
		Version += ` (bindata)`
	}
	config.SetVersion(Version)

	// Service
	if len(os.Args) > 1 && !strings.HasPrefix(os.Args[1], `-`) {
		config.MustOK(config.InitConfig())
		application.WatchConfig(config.InitConfig, false)
		if err := service.Run(os.Args[1]); err != nil {
			stdLog.Println(err)
		}
		return
	}

	err := config.ParseConfig()
	if err != nil {
		if config.IsInstalled() {
			config.MustOK(err)
		} else {
			log.Error(err)
		}
	}

	if config.DefaultCLIConfig.OnlyRunServer() {
		return
	}

	//Manager
	config.DefaultCLIConfig.RunStartup()

	if config.IsInstalled() {
		// 继续上次任务
		if err := cron.InitJobs(); err != nil {
			log.Error(err)
		}
	}

	e := echo.New()
	if binData {
		e.SetDebug(false)
		log.SetLevel(`Info`)
	} else {
		e.SetDebug(true)
		log.SetLevel(`Debug`)
	}
	e.Use(middleware.Log(), middleware.Recover())
	e.Use(middleware.Gzip(&middleware.GzipConfig{
		Skipper: func(c echo.Context) bool {
			switch c.Request().URL().Path() {
			case `/manage/cmdSend/info`, `/download/progress/info`:
				return true
			}
			return false
		},
	}))
	e.Use(func(h echo.Handler) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Response().Header().Set(`Server`, `nging/`+Version)
			return h.Handle(c)
		}
	})

	// 注册静态资源文件(网站素材文件)
	e.Use(staticMW)

	// 启用session
	e.Use(session.Middleware(config.SessionOptions))

	// 启用多语言支持
	config.DefaultConfig.Language.SetFSFunc(langFSFunc)
	e.Use(language.New(&config.DefaultConfig.Language).Middleware())

	// 启用Validation
	e.Use(middleware.Validate(echo.NewValidation))

	// 注册模板引擎
	renderOptions := &render.Config{
		TmplDir: `./template`,
		Engine:  `standard`,
		ParseStrings: map[string]string{
			`__PUBLIC__`: `/public`,
			`__ASSETS__`: `/public/assets`,
			`__TMPL__`:   `./template`,
		},
		Reload:     true,
		ErrorPages: config.DefaultConfig.Sys.ErrorPages,
	}
	renderOptions.ApplyTo(e)
	if tmplMgr != nil {
		renderOptions.Renderer().SetManager(tmplMgr)
	}
	events.AddEvent(`clearCache`, func(next func(r bool), args ...interface{}) {
		renderOptions.Renderer().ClearCache()
		next(true)
	})

	application.Initialize(e)
	c := &engine.Config{
		Address:     fmt.Sprintf(`:%v`, config.DefaultCLIConfig.Port),
		TLSAuto:     false,
		TLSCacheDir: config.DefaultConfig.Sys.SSLCacheDir,
		TLSCertFile: config.DefaultConfig.Sys.SSLCertFile,
		TLSKeyFile:  config.DefaultConfig.Sys.SSLKeyFile,
	}
	e.Run(standard.NewWithConfig(c))
}
<commit_msg>update<commit_after>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

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

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

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

*/
package main

import (
	"flag"
	"fmt"
	stdLog "log"
	"net/http"
	"os"
	"strings"

	_ "github.com/admpub/nging/application/library/sqlite"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/engine"
	"github.com/webx-top/echo/engine/standard"
	"github.com/webx-top/echo/handler/mvc/events"
	"github.com/webx-top/echo/middleware"
	"github.com/webx-top/echo/middleware/language"
	"github.com/webx-top/echo/middleware/render"
	"github.com/webx-top/echo/middleware/render/driver"
	"github.com/webx-top/echo/middleware/session"

	"github.com/admpub/log"
	"github.com/admpub/nging/application"
	"github.com/admpub/nging/application/library/config"
	"github.com/admpub/nging/application/library/cron"
	"github.com/admpub/nging/application/library/service"
)

var (
	BUILD_TIME string
	CLOUD_GOX  string
	COMMIT     string
	VERSION    = `1.2.0`

	version    string
	binData    bool
	staticMW   interface{}
	tmplMgr    driver.Manager
	langFSFunc func(dir string) http.FileSystem
)

func main() {
	config.DefaultCLIConfig.InitFlag()
	flag.Parse()
	if binData {
		version = VERSION + ` (bindata)`
	} else {
		version = VERSION
	}
	config.SetVersion(version)

	// Service
	if len(os.Args) > 1 && !strings.HasPrefix(os.Args[1], `-`) {
		config.MustOK(config.InitConfig())
		application.WatchConfig(config.InitConfig, false)
		if err := service.Run(os.Args[1]); err != nil {
			stdLog.Println(err)
		}
		return
	}

	err := config.ParseConfig()
	if err != nil {
		if config.IsInstalled() {
			config.MustOK(err)
		} else {
			log.Error(err)
		}
	}

	if config.DefaultCLIConfig.OnlyRunServer() {
		return
	}

	//Manager
	config.DefaultCLIConfig.RunStartup()

	if config.IsInstalled() {
		// 继续上次任务
		if err := cron.InitJobs(); err != nil {
			log.Error(err)
		}
	}

	e := echo.New()
	if binData {
		e.SetDebug(false)
		log.SetLevel(`Info`)
	} else {
		e.SetDebug(true)
		log.SetLevel(`Debug`)
	}
	e.Use(middleware.Log(), middleware.Recover())
	e.Use(middleware.Gzip(&middleware.GzipConfig{
		Skipper: func(c echo.Context) bool {
			switch c.Request().URL().Path() {
			case `/manage/cmdSend/info`, `/download/progress/info`:
				return true
			}
			return false
		},
	}))
	e.Use(func(h echo.Handler) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Response().Header().Set(`Server`, `nging/`+Version)
			return h.Handle(c)
		}
	})

	// 注册静态资源文件(网站素材文件)
	e.Use(staticMW)

	// 启用session
	e.Use(session.Middleware(config.SessionOptions))

	// 启用多语言支持
	config.DefaultConfig.Language.SetFSFunc(langFSFunc)
	e.Use(language.New(&config.DefaultConfig.Language).Middleware())

	// 启用Validation
	e.Use(middleware.Validate(echo.NewValidation))

	// 注册模板引擎
	renderOptions := &render.Config{
		TmplDir: `./template`,
		Engine:  `standard`,
		ParseStrings: map[string]string{
			`__PUBLIC__`: `/public`,
			`__ASSETS__`: `/public/assets`,
			`__TMPL__`:   `./template`,
		},
		Reload:     true,
		ErrorPages: config.DefaultConfig.Sys.ErrorPages,
	}
	renderOptions.ApplyTo(e)
	if tmplMgr != nil {
		renderOptions.Renderer().SetManager(tmplMgr)
	}
	events.AddEvent(`clearCache`, func(next func(r bool), args ...interface{}) {
		renderOptions.Renderer().ClearCache()
		next(true)
	})

	application.Initialize(e)
	c := &engine.Config{
		Address:     fmt.Sprintf(`:%v`, config.DefaultCLIConfig.Port),
		TLSAuto:     false,
		TLSCacheDir: config.DefaultConfig.Sys.SSLCacheDir,
		TLSCertFile: config.DefaultConfig.Sys.SSLCertFile,
		TLSKeyFile:  config.DefaultConfig.Sys.SSLKeyFile,
	}
	e.Run(standard.NewWithConfig(c))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/russross/blackfriday"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"strconv"
	"text/template"
	htmltemplate "html/template"
	"regexp"
)

var h1RegEx = regexp.MustCompile("<[\\w]*h1.*>(.*)</[\\w]*h1[\\w]*>")

func main() {
	var docsDirPath string
	var port int

	flag.StringVar(&docsDirPath, "docsdir", "", "Root Directory for all the Docs")
	flag.IntVar(&port, "port", 9000, "Port to run the server on")
	flag.Parse()

	if docsDirPath == "" {
		flag.PrintDefaults()
		return
	}

	var err error
	docsDir, err := os.Open(docsDirPath) // For read access.
	if err != nil {
		log.Fatal(err)
		return
	}

	fi, err := docsDir.Stat()
	if err != nil {
		log.Fatal(err)
		return
	}

	if !fi.IsDir() {
		log.Fatal("%v is not a directory.", docsDirPath)
		return
	}

	editServer := http.FileServer(http.Dir("editor"))
	fileServer := http.FileServer(http.Dir(docsDirPath))

 	http.Handle("/editor/", http.StripPrefix("/editor/",editServer))
	http.Handle("/", &MarkdownHandler{FileRoot: docsDir.Name(), FileServer: fileServer, EditServer: editServer})

	log.Fatal(http.ListenAndServe(":" + strconv.Itoa(port), nil))
}

type MarkdownHandler struct {
	FileServer http.Handler
	EditServer http.Handler
	FileRoot   string
}

func (h *MarkdownHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	rp := r.URL.Path

	var hasMdFile = false
	var filename, mdUrl string
	if strings.HasSuffix(rp, ".html") {
		mdUrl = strings.TrimSuffix(rp, ".html") + ".md"
		filename = filepath.Join(h.FileRoot, mdUrl)

		if _, err := os.Stat(filename); err == nil {
			hasMdFile = true;
		}
	}

	if !hasMdFile {
		filename = filepath.Join(h.FileRoot, rp)
	}

	if r.Method == "POST" {
		if err := os.MkdirAll(filepath.Dir(filename), 0744); err != nil {
			log.Printf("Error Making Directories: %v %v", filename, err)
			return
		}

		body, err := ioutil.ReadAll(r.Body)
		if err == nil {
			ioutil.WriteFile(filename, body, 0644)
		}
		return
	}

	r.ParseForm()
	_, editMode := r.Form["edit"]

	if editMode {
		if hasMdFile {
			w.Header()["Location"] = []string{fmt.Sprintf("%s?%s", mdUrl, r.URL.RawQuery)}
			w.WriteHeader(http.StatusFound)
		}
		w.Header()["Content-Type"] = []string{"text/html"}

		t, _ := htmltemplate.ParseFiles("editor/editor.html")

		content, _ := ioutil.ReadFile(filename)

		data := make(map[string]string)
		data["path"] = rp;
		data["content"] = string(content)

		t.Execute(w, data)
		return
	}

	if !hasMdFile {
		h.FileServer.ServeHTTP(w, r)
		return
	}

	w.Header()["Content-Type"] = []string{"text/html"}

	flags := blackfriday.HTML_TOC | blackfriday.HTML_GITHUB_BLOCKCODE

	extensions := 0
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS

	body, _ := ioutil.ReadFile(filename)
	output := blackfriday.Markdown(body, blackfriday.HtmlRenderer( flags,"",""), extensions)

	navClosingTag := []byte{'<','/','n','a','v','>'}
	navMarker := bytes.Index(output, navClosingTag)
	navMarker = navMarker + len(navClosingTag)

	toc := output[:navMarker]
	content := output[navMarker:]

	var title []byte
	if headingSearch := h1RegEx.FindSubmatch(content); headingSearch != nil && len(headingSearch) >= 2 {
		title = headingSearch[1]
	}

	data := make(map[string]string)
	data["content"] = string(content)
	data["toc"] = string(toc)
	data["title"] = string(title)

	t, _ := template.ParseFiles(filepath.Join(h.FileRoot, "template.html"))

	t.Execute(w, data)
	return
}
<commit_msg>Authentication on the editor.<commit_after>package main

import (
	"code.google.com/p/go.crypto/bcrypt"
	"bufio"
	"bytes"
	"errors"
	"encoding/base64"
	"encoding/gob"
	"flag"
	"fmt"
	"github.com/howeyc/gopass"
	"github.com/russross/blackfriday"
	htmltemplate "html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"text/template"
)

var h1RegEx = regexp.MustCompile("<[\\w]*h1.*>(.*)</[\\w]*h1[\\w]*>")

func main() {
	var home, relDocsDirPath, relEditorDirPath string
	var port int
	var isAddUser bool

	flag.StringVar(&home, "home", "/opt/docserver", "The root directory from which all other paths are relative.")
	flag.StringVar(&relDocsDirPath, "docsdir", "content", "Directory that contains all the documentation.")
	flag.StringVar(&relEditorDirPath, "editor", "editor", "Directory that contains the editor files.")
	flag.IntVar(&port, "port", 9000, "The port to bind the server to.")
	flag.BoolVar(&isAddUser, "adduser", false, "Instead of running the doc server, add a user to the password file.")
	flag.Parse()

	usersFile := filepath.Join(home, "users")

	users, err := decodeUserFile(usersFile)
	if err != nil {
		if !isAddUser {
			log.Println("Failed to open users file. Either it doesn't exist or cannot be accessed. Use the -adduser flag to create a new one.")
			return
		}
		users = make(map[string][]byte)
	}

	if isAddUser {
		addUser(usersFile, users)
		return
	}

	if home == "" || relDocsDirPath == "" || relEditorDirPath == "" || port <= 0 || port > 65535 {
		flag.PrintDefaults()
		return
	}

	unsanitisedDocsDirPath := filepath.Join(home, relDocsDirPath)
	docsDirPath, err := sanitisePath(unsanitisedDocsDirPath)
	if err != nil {
		log.Println(fmt.Sprintf("Failed to open docs directory for reading: %v, %v", docsDirPath, err))
		return
	}

	unsanitiedEditorDirPath := filepath.Join(home, relEditorDirPath)
	editorDirPath, err := sanitisePath(unsanitiedEditorDirPath)
	if err != nil {
		log.Println(fmt.Sprintf("Failed to open editor directory for reading: %v, %v", editorDirPath, err))
		return
	}

	editServer := http.FileServer(http.Dir(editorDirPath))
	fileServer := http.FileServer(http.Dir(docsDirPath))

	http.Handle("/editor/", http.StripPrefix("/editor/", editServer))
	http.Handle("/", &MarkdownHandler{
		DocRoot:    docsDirPath,
		EditorRoot: editorDirPath,
		FileServer: fileServer,
		EditServer: editServer,
		Users: users})

	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(port), nil))
}

type MarkdownHandler struct {
	FileServer http.Handler
	EditServer http.Handler
	DocRoot    string
	EditorRoot string
	Users      map[string][]byte
}

func (h *MarkdownHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	rp := r.URL.Path

	var hasMdFile = false
	var filename, mdUrl string
	if strings.HasSuffix(rp, ".html") {
		mdUrl = strings.TrimSuffix(rp, ".html") + ".md"
		filename = filepath.Join(h.DocRoot, mdUrl)

		if _, err := os.Stat(filename); err == nil {
			hasMdFile = true
		}
	}

	if !hasMdFile {
		filename = filepath.Join(h.DocRoot, rp)
	}

	if r.Method == "POST" {
		if requiresAuth(w, r, h.Users) {
			return
		}

		if err := os.MkdirAll(filepath.Dir(filename), 0744); err != nil {
			log.Printf("Error making directories: %v %v", filename, err)
			return
		}

		body, err := ioutil.ReadAll(r.Body)
		if err == nil {
			ioutil.WriteFile(filename, body, 0644)
		}
		return
	}

	r.ParseForm()
	_, editMode := r.Form["edit"]

	if editMode {
		if requiresAuth(w, r, h.Users) {
			return
		}

		if hasMdFile {
			w.Header()["Location"] = []string{fmt.Sprintf("%s?%s", mdUrl, r.URL.RawQuery)}
			w.WriteHeader(http.StatusFound)
		}
		w.Header()["Content-Type"] = []string{"text/html"}

		editorHtml := filepath.Join(h.EditorRoot, "editor.html")
		t, _ := htmltemplate.ParseFiles(editorHtml)

		content, _ := ioutil.ReadFile(filename)

		data := make(map[string]string)
		data["path"] = rp
		data["content"] = string(content)

		t.Execute(w, data)
		return
	}

	if !hasMdFile {
		h.FileServer.ServeHTTP(w, r)
		return
	}

	w.Header()["Content-Type"] = []string{"text/html"}

	flags := blackfriday.HTML_TOC | blackfriday.HTML_GITHUB_BLOCKCODE

	extensions := 0
	extensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	extensions |= blackfriday.EXTENSION_TABLES
	extensions |= blackfriday.EXTENSION_FENCED_CODE
	extensions |= blackfriday.EXTENSION_AUTOLINK
	extensions |= blackfriday.EXTENSION_STRIKETHROUGH
	extensions |= blackfriday.EXTENSION_SPACE_HEADERS

	body, _ := ioutil.ReadFile(filename)
	output := blackfriday.Markdown(body, blackfriday.HtmlRenderer(flags, "", ""), extensions)

	navClosingTag := []byte{'<', '/', 'n', 'a', 'v', '>'}
	navMarker := bytes.Index(output, navClosingTag)
	navMarker = navMarker + len(navClosingTag)

	toc := output[:navMarker]
	content := output[navMarker:]

	var title []byte
	if headingSearch := h1RegEx.FindSubmatch(content); headingSearch != nil && len(headingSearch) >= 2 {
		title = headingSearch[1]
	}

	data := make(map[string]string)
	data["content"] = string(content)
	data["toc"] = string(toc)
	data["title"] = string(title)

	t, _ := template.ParseFiles(filepath.Join(h.DocRoot, "template.html"))

	t.Execute(w, data)
	return
}

func sanitisePath(path string) (string, error) {
	directory, err := os.Open(path)
	defer directory.Close()
	if err != nil {
		return "", err
	}

	fi, err := directory.Stat()
	if err != nil {
		return "", err
	}

	if !fi.IsDir() {
		return "", errors.New("Path is not directory")
	}

	return directory.Name(), nil
}

func addUser(path string, usersMap map[string][]byte) {
	bio := bufio.NewReader(os.Stdin)

	fmt.Printf("Username: ")
	username, _ := bio.ReadString('\n')
	username = strings.TrimSpace(username)
	fmt.Printf("Password: ")
    password := string(gopass.GetPasswdMasked())

    passwordHash, err := bcrypt.GenerateFromPassword([]byte(password), 0)
    if err != nil {
    	panic(err)
    }

    log.Printf("Saving Hash: '%v'", passwordHash)

    usersMap[username] = passwordHash

    encodeUserFile(path, usersMap)
}

func decodeUserFile(path string) (map[string][]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	decoder := gob.NewDecoder(file)

	var users = make(map[string][]byte)

	if err := decoder.Decode(&users); err != nil {
		return nil, err
	}

	return users, nil
}

func encodeUserFile(path string, users map[string][]byte) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	encoder := gob.NewEncoder(file)

	if err := encoder.Encode(users); err != nil {
		log.Printf("err: %v", err)
		return err
	}

	return nil
}

func requiresAuth(w http.ResponseWriter, r *http.Request, users map[string][]byte) bool {
	authHeader := r.Header.Get("Authorization");

	authHeaderParts := strings.SplitN(authHeader, " ", 2)
	if len(authHeaderParts) != 2 || authHeaderParts[0] != "Basic" {
		sendAuthHeaders(w)
		return true
	}

	decoded, err := base64.StdEncoding.DecodeString(authHeaderParts[1])
	if err != nil {
		sendAuthHeaders(w)
		return true
	}

	credentials := strings.SplitN(string(decoded), ":", 2)
	if len(credentials) != 2 {
		sendAuthHeaders(w)
		return true
	}

	username := credentials[0]
	password := credentials[1]
	storedHash := users[username]

	log.Printf("Username: '%v'", username)
	log.Printf("Password: '%v'", password)
	log.Printf("Stored Hash: '%v'", storedHash)

	if err := bcrypt.CompareHashAndPassword(storedHash, []byte(password)); err != nil {
		log.Printf("Error: %v", err)
		sendAuthHeaders(w)
		return true
	}

	return false;
}

func sendAuthHeaders(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", `Basic realm="editdocs"`)
	w.WriteHeader(401)
	w.Write([]byte("401 Unauthorized\n"))
}<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"html"
	"io"
	"io/ioutil"
	"log"
	"net/http"

	"github.com/gorilla/mux"

	"github.com/dchest/uniuri"
	"github.com/ewhal/pygments"
	_ "github.com/mattn/go-sqlite3"
)

const (
	ADDRESS = "http://localhost:9900"
	LENGTH  = 6
	TEXT    = "$ <command> | curl -F 'p=<-' " + ADDRESS + "\n"
	PORT    = ":9900"
)

func check(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

func generateName() string {
	s := uniuri.NewLen(LENGTH)
	db, err := sql.Open("sqlite3", "./database.db")
	check(err)

	query, err := db.Query("select id from pastebin")
	for query.Next() {
		var id string
		err := query.Scan(&id)
		if err != nil {

		}
		if id == s {
			generateName()
		}
	}
	db.Close()

	return s

}
func save(raw []byte) string {
	paste := raw[86 : len(raw)-46]

	s := generateName()
	db, err := sql.Open("sqlite3", "./database.db")
	check(err)
	stmt, err := db.Prepare("INSERT INTO pastebin(id, hash, data, delkey) values(?,?,?,?)")
	_, err = stmt.Exec(s, html.EscapeString(string(paste)))
	check(err)
	db.Close()

	return s
}

func saveHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "POST":
		buf, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), 500)
		}
		name := save(buf)
		io.WriteString(w, ADDRESS+"/p/"+name+"\n")
	}

}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, TEXT)
}
func langHandler(w http.ResponseWriter, r *http.Request) {

	vars := mux.Vars(r)
	paste := vars["pasteId"]
	lang := vars["lang"]
	s := getPaste(paste)
	highlight := pygments.Highlight(html.UnescapeString(s), lang, "html", "full, style=autumn,linenos=True, lineanchors=True,anchorlinenos=True,", "utf-8")
	io.WriteString(w, highlight)

}

func getPaste(paste string) string {
	param1 := html.EscapeString(paste)
	db, err := sql.Open("sqlite3", "./database.db")
	var s string
	err = db.QueryRow("select data from pastebin where id=?", param1).Scan(&s)
	db.Close()
	check(err)

	if err == sql.ErrNoRows {
		return "Error invalid paste"
	} else {
		return html.UnescapeString(s)
	}

}
func pasteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s := getPaste(paste)
	io.WriteString(w, s)

}

func main() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/", rootHandler)
	router.HandleFunc("/p/{pasteId}", pasteHandler)
	router.HandleFunc("/p/{pasteId}/{lang}", langHandler)
	router.HandleFunc("/save", saveHandler)
	router.HandleFunc("/save", saveHandler)
	err := http.ListenAndServe(PORT, router)
	if err != nil {
		log.Fatal(err)
	}

}
<commit_msg>Add json output<commit_after>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"html"
	"io"
	"io/ioutil"
	"log"
	"net/http"

	"github.com/dchest/uniuri"
	"github.com/ewhal/pygments"
	"github.com/gorilla/mux"
	_ "github.com/mattn/go-sqlite3"
)

const (
	ADDRESS = "http://localhost:9900"
	LENGTH  = 6
	TEXT    = "$ <command> | curl -F 'p=<-' " + ADDRESS + "\n"
	PORT    = ":9900"
)

type Response struct {
	ID     string `json:"id"`
	HASH   string `json:"hash"`
	URL    string `json:"url"`
	SIZE   int    `json:"size"`
	DELKEY string `json:"delkey"`
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

func generateName() string {
	s := uniuri.NewLen(LENGTH)
	db, err := sql.Open("sqlite3", "./database.db")
	check(err)

	query, err := db.Query("select id from pastebin")
	for query.Next() {
		var id string
		err := query.Scan(&id)
		if err != nil {

		}
		if id == s {
			generateName()
		}
	}
	db.Close()

	return s

}
func hash(paste []byte) string {
	hasher := sha1.New()
	hasher.Write(paste)
	sha := base64.URLEncoding.EncodeToString(hasher.Sum(nil))
	return sha
}

func save(raw []byte) []string {
	p := raw[86 : len(raw)-46]

	db, err := sql.Open("sqlite3", "./database.db")
	check(err)

	sha := hash(p)
	id := generateName()
	url := ADDRESS + "/p/" + id
	delKey := uniuri.NewLen(40)
	paste := html.EscapeString(string(p))

	stmt, err := db.Prepare("INSERT INTO pastebin(id, hash, data, delkey) values(?,?,?,?)")
	check(err)
	_, err = stmt.Exec(id, sha, paste, delKey)
	check(err)
	db.Close()
	return []string{id, sha, url, paste, delKey}

}

func delHandler(w http.ResponseWriter, r *http.Request) {
}
func saveHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	output := vars["output"]
	switch r.Method {
	case "POST":
		buf, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), 500)
		}

		values := save(buf)

		switch output {
		case "json":
			b := &Response{
				ID:     values[0],
				HASH:   values[1],
				URL:    values[2],
				SIZE:   len(values[3]),
				DELKEY: values[4],
			}

			w.Header().Set("Content-Type", "application/json")
			err := json.NewEncoder(w).Encode(b)

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

		default:
			io.WriteString(w, values[2]+"\n")
		}
	}

}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, TEXT)
}
func langHandler(w http.ResponseWriter, r *http.Request) {

	vars := mux.Vars(r)
	paste := vars["pasteId"]
	lang := vars["lang"]
	s := getPaste(paste)
	highlight := pygments.Highlight(html.UnescapeString(s), lang, "html", "full, style=autumn,linenos=True, lineanchors=True,anchorlinenos=True,", "utf-8")
	io.WriteString(w, highlight)

}

func getPaste(paste string) string {
	param1 := html.EscapeString(paste)
	db, err := sql.Open("sqlite3", "./database.db")
	var s string
	err = db.QueryRow("select data from pastebin where id=?", param1).Scan(&s)
	db.Close()
	check(err)

	if err == sql.ErrNoRows {
		return "Error invalid paste"
	} else {
		return html.UnescapeString(s)
	}

}
func pasteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s := getPaste(paste)
	io.WriteString(w, s)

}

func main() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/", rootHandler)
	router.HandleFunc("/p/{pasteId}", pasteHandler)
	router.HandleFunc("/p/{pasteId}/{lang}", langHandler)
	router.HandleFunc("/save", saveHandler)
	router.HandleFunc("/save/{output}", saveHandler)
	router.HandleFunc("/del/{pasteId}/{delKey}", delHandler)
	err := http.ListenAndServe(PORT, router)
	if err != nil {
		log.Fatal(err)
	}

}
<|endoftext|>


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

// sshtunnel is daemon for setting up forward and reverse SSH tunnels.
//
// The daemon is started by executing sshtunnel with the path to a JSON
// configuration file. The configuration takes the following form:
//
//	{
//		"KeyFiles": ["/path/to/key.priv"],
//		"KnownHostFiles": ["/path/to/known_hosts"],
//		"Tunnels": [{
//			// Forward tunnel (locally binded socket proxies to remote target).
//			"Tunnel": "bind_address:port -> dial_address:port",
//			"Server": "user@host:port",
//		}, {
//			// Reverse tunnel (remotely binded socket proxies to local target).
//			"Tunnel": "dial_address:port <- bind_address:port",
//			"Server": "user@host:port",
//		}],
//	}
//
// See the TunnelConfig struct for more details.
package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"os/user"
	"path"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/dsnet/golib/jsonfmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

// Version of the sshtunnel binary. May be set by linker when building.
var version string

type TunnelConfig struct {
	// LogFile is where the proxy daemon will direct its output log.
	// If the path is empty, then the server will output to os.Stderr.
	LogFile string `json:",omitempty"`

	// KeyFiles is a list of SSH private key files.
	KeyFiles []string

	// KnownHostFiles is a list of key database files for host public keys
	// in the OpenSSH known_hosts file format.
	KnownHostFiles []string

	// KeepAlive sets the keep alive settings for each SSH connection.
	// It is recommended that these values match the AliveInterval and
	// AliveCountMax parameters on the remote OpenSSH server.
	// If unset, then the default is an interval of 30s with 2 max counts.
	KeepAlive *KeepAliveConfig `json:",omitempty"`

	// Tunnels is a list of tunnels to establish.
	// The same set of SSH keys will be used to authenticate the
	// SSH connection for each server.
	Tunnels []struct {
		// Tunnel is a pair of host:port endpoints that can be configured
		// to either operate as a forward tunnel or a reverse tunnel.
		//
		// The syntax of a forward tunnel is:
		//	"bind_address:port -> dial_address:port"
		//
		// A forward tunnel opens a listening TCP socket on the
		// local side (at bind_address:port) and proxies all traffic to a
		// socket on the remote side (at dial_address:port).
		//
		// The syntax of a reverse tunnel is:
		//	"dial_address:port <- bind_address:port"
		//
		// A reverse tunnel opens a listening TCP socket on the
		// remote side (at bind_address:port) and proxies all traffic to a
		// socket on the local side (at dial_address:port).
		Tunnel string

		// Server is a remote SSH host. It has the following syntax:
		//	"user@host:port"
		//
		// If the user is missing, then it defaults to the current process user.
		// If the port is missing, then it defaults to 22.
		Server string

		// KeepAlive is a tunnel-specific setting of the global KeepAlive.
		// If unspecified, it uses the global KeepAlive settings.
		KeepAlive *KeepAliveConfig `json:",omitempty"`
	}
}

type KeepAliveConfig struct {
	// Interval is the amount of time in seconds to wait before the
	// tunnel client will send a keep-alive message to ensure some minimum
	// traffic on the SSH connection.
	Interval uint

	// CountMax is the maximum number of consecutive failed responses to
	// keep-alive messages the client is willing to tolerate before considering
	// the SSH connection as dead.
	CountMax uint
}

func loadConfig(conf string) (tunns []tunnel, logger *log.Logger, closer func() error) {
	var logBuf bytes.Buffer
	logger = log.New(io.MultiWriter(os.Stderr, &logBuf), "", log.Ldate|log.Ltime|log.Lshortfile)

	var hash string
	if b, _ := ioutil.ReadFile(os.Args[0]); len(b) > 0 {
		hash = fmt.Sprintf("%x", sha256.Sum256(b))
	}

	// Load configuration file.
	var config TunnelConfig
	c, err := ioutil.ReadFile(conf)
	if err != nil {
		logger.Fatalf("unable to read config: %v", err)
	}
	c, _ = jsonfmt.Format(c, jsonfmt.Minify())
	if err := json.Unmarshal(c, &config); err != nil {
		logger.Fatalf("unable to decode config: %v", err)
	}
	for _, t := range config.Tunnels {
		if config.KeepAlive == nil && t.KeepAlive == nil {
			config.KeepAlive = &KeepAliveConfig{Interval: 30, CountMax: 2}
			break
		}
	}

	// Print the configuration.
	var b bytes.Buffer
	enc := json.NewEncoder(&b)
	enc.SetEscapeHTML(false)
	enc.SetIndent("", "\t")
	enc.Encode(struct {
		TunnelConfig
		BinaryVersion string `json:",omitempty"`
		BinarySHA256  string `json:",omitempty"`
	}{config, version, hash})
	logger.Printf("loaded config:\n%s", b.String())

	// Setup the log output.
	if config.LogFile == "" {
		logger.SetOutput(os.Stderr)
		closer = func() error { return nil }
	} else {
		f, err := os.OpenFile(config.LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0664)
		if err != nil {
			logger.Fatalf("error opening log file: %v", err)
		}
		f.Write(logBuf.Bytes()) // Write log output prior to this point
		logger.Printf("suppress stderr logging (redirected to %s)", f.Name())
		logger.SetOutput(f)
		closer = f.Close
	}

	// Parse all of the private keys.
	var keys []ssh.Signer
	if len(config.KeyFiles) == 0 {
		logger.Fatal("no private keys specified")
	}
	for _, kf := range config.KeyFiles {
		b, err := ioutil.ReadFile(kf)
		if err != nil {
			logger.Fatalf("private key error: %v", err)
		}
		k, err := ssh.ParsePrivateKey(b)
		if err != nil {
			logger.Fatalf("private key error: %v", err)
		}
		keys = append(keys, k)
	}
	auth := []ssh.AuthMethod{ssh.PublicKeys(keys...)}

	// Parse all of the host public keys.
	if len(config.KnownHostFiles) == 0 {
		logger.Fatal("no host public keys specified")
	}
	hostKeys, err := knownhosts.New(config.KnownHostFiles...)
	if err != nil {
		logger.Fatalf("public key error: %v", err)
	}

	// Parse all of the tunnels.
	for _, t := range config.Tunnels {
		var tunn tunnel
		tt := strings.Fields(t.Tunnel)
		if len(tt) != 3 {
			logger.Fatalf("invalid tunnel syntax: %s", t.Tunnel)
		}

		// Parse for the tunnel endpoints.
		switch tt[1] {
		case "->":
			tunn.bindAddr, tunn.mode, tunn.dialAddr = tt[0], '>', tt[2]
		case "<-":
			tunn.dialAddr, tunn.mode, tunn.bindAddr = tt[0], '<', tt[2]
		default:
			logger.Fatalf("invalid tunnel syntax: %s", t.Tunnel)
		}
		for _, addr := range []string{tunn.bindAddr, tunn.dialAddr} {
			if _, _, err := net.SplitHostPort(addr); err != nil {
				logger.Fatalf("invalid endpoint: %s", addr)
			}
		}

		// Parse for the SSH target host.
		tunn.hostAddr = t.Server
		if i := strings.IndexByte(t.Server, '@'); i >= 0 {
			tunn.user = t.Server[:i]
			tunn.hostAddr = t.Server[i+1:]
		}
		if _, _, err := net.SplitHostPort(tunn.hostAddr); err != nil {
			tunn.hostAddr = net.JoinHostPort(tunn.hostAddr, "22")
		}
		if _, _, err := net.SplitHostPort(tunn.hostAddr); err != nil {
			logger.Fatalf("invalid server: %s", t.Server)
		}

		// Parse for the SSH user.
		if tunn.user == "" {
			u, err := user.Current()
			if err != nil {
				logger.Fatalf("unexpected error: %v", err)
			}
			tunn.user = u.Username
		}

		if t.KeepAlive == nil {
			tunn.keepAlive = *config.KeepAlive
		} else {
			tunn.keepAlive = *t.KeepAlive
		}
		tunn.retryInterval = 30 * time.Second
		tunn.auth = auth
		tunn.hostKeys = hostKeys
		tunn.log = logger
		tunns = append(tunns, tunn)
	}

	return tunns, logger, closer
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "Usage:\n")
		fmt.Fprintf(os.Stderr, "\t%s CONFIG_PATH\n", os.Args[0])
		os.Exit(1)
	}
	tunns, logger, closer := loadConfig(os.Args[1])
	defer closer()

	// Setup signal handler to initiate shutdown.
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		sigc := make(chan os.Signal, 1)
		signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM)
		logger.Printf("received %v - initiating shutdown", <-sigc)
		cancel()
	}()

	// Start a bridge for each tunnel.
	var wg sync.WaitGroup
	logger.Printf("%s starting", path.Base(os.Args[0]))
	defer logger.Printf("%s shutdown", path.Base(os.Args[0]))
	for _, t := range tunns {
		wg.Add(1)
		go t.bindTunnel(ctx, &wg)
	}
	wg.Wait()
}
<commit_msg>Log jsonfmt.Format error<commit_after>// Copyright 2016, The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

// sshtunnel is daemon for setting up forward and reverse SSH tunnels.
//
// The daemon is started by executing sshtunnel with the path to a JSON
// configuration file. The configuration takes the following form:
//
//	{
//		"KeyFiles": ["/path/to/key.priv"],
//		"KnownHostFiles": ["/path/to/known_hosts"],
//		"Tunnels": [{
//			// Forward tunnel (locally binded socket proxies to remote target).
//			"Tunnel": "bind_address:port -> dial_address:port",
//			"Server": "user@host:port",
//		}, {
//			// Reverse tunnel (remotely binded socket proxies to local target).
//			"Tunnel": "dial_address:port <- bind_address:port",
//			"Server": "user@host:port",
//		}],
//	}
//
// See the TunnelConfig struct for more details.
package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"os/user"
	"path"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/dsnet/golib/jsonfmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

// Version of the sshtunnel binary. May be set by linker when building.
var version string

type TunnelConfig struct {
	// LogFile is where the proxy daemon will direct its output log.
	// If the path is empty, then the server will output to os.Stderr.
	LogFile string `json:",omitempty"`

	// KeyFiles is a list of SSH private key files.
	KeyFiles []string

	// KnownHostFiles is a list of key database files for host public keys
	// in the OpenSSH known_hosts file format.
	KnownHostFiles []string

	// KeepAlive sets the keep alive settings for each SSH connection.
	// It is recommended that these values match the AliveInterval and
	// AliveCountMax parameters on the remote OpenSSH server.
	// If unset, then the default is an interval of 30s with 2 max counts.
	KeepAlive *KeepAliveConfig `json:",omitempty"`

	// Tunnels is a list of tunnels to establish.
	// The same set of SSH keys will be used to authenticate the
	// SSH connection for each server.
	Tunnels []struct {
		// Tunnel is a pair of host:port endpoints that can be configured
		// to either operate as a forward tunnel or a reverse tunnel.
		//
		// The syntax of a forward tunnel is:
		//	"bind_address:port -> dial_address:port"
		//
		// A forward tunnel opens a listening TCP socket on the
		// local side (at bind_address:port) and proxies all traffic to a
		// socket on the remote side (at dial_address:port).
		//
		// The syntax of a reverse tunnel is:
		//	"dial_address:port <- bind_address:port"
		//
		// A reverse tunnel opens a listening TCP socket on the
		// remote side (at bind_address:port) and proxies all traffic to a
		// socket on the local side (at dial_address:port).
		Tunnel string

		// Server is a remote SSH host. It has the following syntax:
		//	"user@host:port"
		//
		// If the user is missing, then it defaults to the current process user.
		// If the port is missing, then it defaults to 22.
		Server string

		// KeepAlive is a tunnel-specific setting of the global KeepAlive.
		// If unspecified, it uses the global KeepAlive settings.
		KeepAlive *KeepAliveConfig `json:",omitempty"`
	}
}

type KeepAliveConfig struct {
	// Interval is the amount of time in seconds to wait before the
	// tunnel client will send a keep-alive message to ensure some minimum
	// traffic on the SSH connection.
	Interval uint

	// CountMax is the maximum number of consecutive failed responses to
	// keep-alive messages the client is willing to tolerate before considering
	// the SSH connection as dead.
	CountMax uint
}

func loadConfig(conf string) (tunns []tunnel, logger *log.Logger, closer func() error) {
	var logBuf bytes.Buffer
	logger = log.New(io.MultiWriter(os.Stderr, &logBuf), "", log.Ldate|log.Ltime|log.Lshortfile)

	var hash string
	if b, _ := ioutil.ReadFile(os.Args[0]); len(b) > 0 {
		hash = fmt.Sprintf("%x", sha256.Sum256(b))
	}

	// Load configuration file.
	var config TunnelConfig
	c, err := ioutil.ReadFile(conf)
	if err != nil {
		logger.Fatalf("unable to read config: %v", err)
	}
	if c, err = jsonfmt.Format(c, jsonfmt.Standardize()); err != nil {
		logger.Fatalf("unable to parse config: %v", err)
	}
	if err := json.Unmarshal(c, &config); err != nil {
		logger.Fatalf("unable to decode config: %v", err)
	}
	for _, t := range config.Tunnels {
		if config.KeepAlive == nil && t.KeepAlive == nil {
			config.KeepAlive = &KeepAliveConfig{Interval: 30, CountMax: 2}
			break
		}
	}

	// Print the configuration.
	var b bytes.Buffer
	enc := json.NewEncoder(&b)
	enc.SetEscapeHTML(false)
	enc.SetIndent("", "\t")
	enc.Encode(struct {
		TunnelConfig
		BinaryVersion string `json:",omitempty"`
		BinarySHA256  string `json:",omitempty"`
	}{config, version, hash})
	logger.Printf("loaded config:\n%s", b.String())

	// Setup the log output.
	if config.LogFile == "" {
		logger.SetOutput(os.Stderr)
		closer = func() error { return nil }
	} else {
		f, err := os.OpenFile(config.LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0664)
		if err != nil {
			logger.Fatalf("error opening log file: %v", err)
		}
		f.Write(logBuf.Bytes()) // Write log output prior to this point
		logger.Printf("suppress stderr logging (redirected to %s)", f.Name())
		logger.SetOutput(f)
		closer = f.Close
	}

	// Parse all of the private keys.
	var keys []ssh.Signer
	if len(config.KeyFiles) == 0 {
		logger.Fatal("no private keys specified")
	}
	for _, kf := range config.KeyFiles {
		b, err := ioutil.ReadFile(kf)
		if err != nil {
			logger.Fatalf("private key error: %v", err)
		}
		k, err := ssh.ParsePrivateKey(b)
		if err != nil {
			logger.Fatalf("private key error: %v", err)
		}
		keys = append(keys, k)
	}
	auth := []ssh.AuthMethod{ssh.PublicKeys(keys...)}

	// Parse all of the host public keys.
	if len(config.KnownHostFiles) == 0 {
		logger.Fatal("no host public keys specified")
	}
	hostKeys, err := knownhosts.New(config.KnownHostFiles...)
	if err != nil {
		logger.Fatalf("public key error: %v", err)
	}

	// Parse all of the tunnels.
	for _, t := range config.Tunnels {
		var tunn tunnel
		tt := strings.Fields(t.Tunnel)
		if len(tt) != 3 {
			logger.Fatalf("invalid tunnel syntax: %s", t.Tunnel)
		}

		// Parse for the tunnel endpoints.
		switch tt[1] {
		case "->":
			tunn.bindAddr, tunn.mode, tunn.dialAddr = tt[0], '>', tt[2]
		case "<-":
			tunn.dialAddr, tunn.mode, tunn.bindAddr = tt[0], '<', tt[2]
		default:
			logger.Fatalf("invalid tunnel syntax: %s", t.Tunnel)
		}
		for _, addr := range []string{tunn.bindAddr, tunn.dialAddr} {
			if _, _, err := net.SplitHostPort(addr); err != nil {
				logger.Fatalf("invalid endpoint: %s", addr)
			}
		}

		// Parse for the SSH target host.
		tunn.hostAddr = t.Server
		if i := strings.IndexByte(t.Server, '@'); i >= 0 {
			tunn.user = t.Server[:i]
			tunn.hostAddr = t.Server[i+1:]
		}
		if _, _, err := net.SplitHostPort(tunn.hostAddr); err != nil {
			tunn.hostAddr = net.JoinHostPort(tunn.hostAddr, "22")
		}
		if _, _, err := net.SplitHostPort(tunn.hostAddr); err != nil {
			logger.Fatalf("invalid server: %s", t.Server)
		}

		// Parse for the SSH user.
		if tunn.user == "" {
			u, err := user.Current()
			if err != nil {
				logger.Fatalf("unexpected error: %v", err)
			}
			tunn.user = u.Username
		}

		if t.KeepAlive == nil {
			tunn.keepAlive = *config.KeepAlive
		} else {
			tunn.keepAlive = *t.KeepAlive
		}
		tunn.retryInterval = 30 * time.Second
		tunn.auth = auth
		tunn.hostKeys = hostKeys
		tunn.log = logger
		tunns = append(tunns, tunn)
	}

	return tunns, logger, closer
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "Usage:\n")
		fmt.Fprintf(os.Stderr, "\t%s CONFIG_PATH\n", os.Args[0])
		os.Exit(1)
	}
	tunns, logger, closer := loadConfig(os.Args[1])
	defer closer()

	// Setup signal handler to initiate shutdown.
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		sigc := make(chan os.Signal, 1)
		signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM)
		logger.Printf("received %v - initiating shutdown", <-sigc)
		cancel()
	}()

	// Start a bridge for each tunnel.
	var wg sync.WaitGroup
	logger.Printf("%s starting", path.Base(os.Args[0]))
	defer logger.Printf("%s shutdown", path.Base(os.Args[0]))
	for _, t := range tunns {
		wg.Add(1)
		go t.bindTunnel(ctx, &wg)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"

	"github.com/apcera/nats"
	cf_debug_server "github.com/cloudfoundry-incubator/cf-debug-server"
	"github.com/cloudfoundry-incubator/cf-lager"
	routing_api "github.com/cloudfoundry-incubator/routing-api"
	token_fetcher "github.com/cloudfoundry-incubator/uaa-token-fetcher"
	"github.com/cloudfoundry/dropsonde"
	"github.com/cloudfoundry/gorouter/access_log"
	vcap "github.com/cloudfoundry/gorouter/common"
	"github.com/cloudfoundry/gorouter/common/secure"
	"github.com/cloudfoundry/gorouter/config"
	"github.com/cloudfoundry/gorouter/proxy"
	rregistry "github.com/cloudfoundry/gorouter/registry"
	"github.com/cloudfoundry/gorouter/route_fetcher"
	"github.com/cloudfoundry/gorouter/router"
	rvarz "github.com/cloudfoundry/gorouter/varz"
	"github.com/cloudfoundry/yagnats"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"

	"flag"
	"fmt"
	"os"
	"runtime"
	"syscall"
	"time"

	"github.com/cloudfoundry/gorouter/metrics"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
)

var configFile string

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func main() {
	flag.StringVar(&configFile, "c", "", "Configuration File")
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	logger, _ := cf_lager.New("gorouter.stdout")
	c := config.DefaultConfig(logger)
	logCounter := vcap.NewLogCounter()

	if configFile != "" {
		c = config.InitConfigFromFile(logger, configFile)
	}

	InitLoggerFromConfig(logger, c, logCounter)
	err := dropsonde.Initialize(c.Logging.MetronAddress, c.Logging.JobName)
	if err != nil {
		logger.Error("Dropsonde failed to initialize ", err)
		os.Exit(1)
	}

	// setup number of procs
	if c.GoMaxProcs != 0 {
		runtime.GOMAXPROCS(c.GoMaxProcs)
	}

	if c.DebugAddr != "" {
		cf_debug_server.Run(c.DebugAddr)
	}

	logger.Info("Setting up NATs connection")
	natsClient := connectToNatsServer(logger, c)

	metricsReporter := metrics.NewMetricsReporter()
	registry := rregistry.NewRouteRegistry(c, natsClient, metricsReporter)

	varz := rvarz.NewVarz(registry)
	compositeReporter := metrics.NewCompositeReporter(varz, metricsReporter)

	accessLogger, err := access_log.CreateRunningAccessLogger(logger, c)
	if err != nil {
		logger.Fatal("Error creating access logger: ", err)
	}

	var crypto secure.Crypto
	var cryptoPrev secure.Crypto
	if c.RouteServiceEnabled {
		crypto = createCrypto(logger, c.RouteServiceSecret)
		if c.RouteServiceSecretPrev != "" {
			cryptoPrev = createCrypto(logger, c.RouteServiceSecretPrev)
		}
	}

	proxy := buildProxy(logger, c, registry, accessLogger, compositeReporter, crypto, cryptoPrev)

	router, err := router.NewRouter(logger, c, proxy, natsClient, registry, varz, logCounter, nil)
	if err != nil {
		logger.Error("An error occurred: ", err)
		os.Exit(1)
	}

	members := grouper.Members{
		{"router", router},
	}
	if c.RoutingApiEnabled() {
		logger.Info("Setting up route fetcher")
		routeFetcher := setupRouteFetcher(logger, c, registry)
		members = append(members, grouper.Member{"router-fetcher", routeFetcher})
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1))

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("gorouter.exited-with-failure: ", err)
		os.Exit(1)
	}

	os.Exit(0)
}

func createCrypto(logger lager.Logger, secret string) *secure.AesGCM {
	// generate secure encryption key using key derivation function (pbkdf2)
	secretPbkdf2 := secure.NewPbkdf2([]byte(secret), 16)
	crypto, err := secure.NewAesGCM(secretPbkdf2)
	if err != nil {
		logger.Error("Error creating route service crypto: %s\n", err)
		os.Exit(1)
	}
	return crypto
}

func buildProxy(logger lager.Logger, c *config.Config, registry rregistry.RegistryInterface, accessLogger access_log.AccessLogger, reporter metrics.ProxyReporter, crypto secure.Crypto, cryptoPrev secure.Crypto) proxy.Proxy {
	args := proxy.ProxyArgs{
		Logger:          logger,
		EndpointTimeout: c.EndpointTimeout,
		Ip:              c.Ip,
		TraceKey:        c.TraceKey,
		Registry:        registry,
		Reporter:        reporter,
		AccessLogger:    accessLogger,
		SecureCookies:   c.SecureCookies,
		TLSConfig: &tls.Config{
			CipherSuites:       c.CipherSuites,
			InsecureSkipVerify: c.SSLSkipValidation,
		},
		RouteServiceEnabled: c.RouteServiceEnabled,
		RouteServiceTimeout: c.RouteServiceTimeout,
		Crypto:              crypto,
		CryptoPrev:          cryptoPrev,
		ExtraHeadersToLog:   c.ExtraHeadersToLog,
	}
	return proxy.NewProxy(args)
}

func setupRouteFetcher(logger lager.Logger, c *config.Config, registry rregistry.RegistryInterface) *route_fetcher.RouteFetcher {
	clock := clock.NewClock()

	tokenFetcher := newTokenFetcher(logger, clock, c)
	_, err := tokenFetcher.FetchToken(false)
	if err != nil {
		logger.Error("Unable to fetch token: ", err)
		os.Exit(1)
	}

	routingApiUri := fmt.Sprintf("%s:%d", c.RoutingApi.Uri, c.RoutingApi.Port)
	routingApiClient := routing_api.NewClient(routingApiUri)

	routeFetcher := route_fetcher.NewRouteFetcher(logger, tokenFetcher, registry, c, routingApiClient, 1, clock)
	return routeFetcher
}

func newTokenFetcher(logger lager.Logger, clock clock.Clock, c *config.Config) token_fetcher.TokenFetcher {
	if c.RoutingApi.AuthDisabled {
		logger.Info("using noop token fetcher")
		return token_fetcher.NewNoOpTokenFetcher()
	}
	tokenFetcherConfig := token_fetcher.TokenFetcherConfig{
		MaxNumberOfRetries:   c.TokenFetcherMaxRetries,
		RetryInterval:        c.TokenFetcherRetryInterval,
		ExpirationBufferTime: c.TokenFetcherExpirationBufferTimeInSeconds,
	}

	tokenFetcher, err := token_fetcher.NewTokenFetcher(logger, &c.OAuth, tokenFetcherConfig, clock)
	if err != nil {
		logger.Error("Error creating token fetcher: %s\n", err)
		os.Exit(1)
	}
	logger.Info("using uaa token fetcher")
	return tokenFetcher
}

func connectToNatsServer(logger lager.Logger, c *config.Config) yagnats.NATSConn {
	var natsClient yagnats.NATSConn
	var err error

	natsServers := c.NatsServers()
	attempts := 3
	for attempts > 0 {
		natsClient, err = yagnats.Connect(natsServers)
		if err == nil {
			break
		} else {
			attempts--
			time.Sleep(100 * time.Millisecond)
		}
	}

	if err != nil {
		logger.Error("Error connecting to NATS: %s\n", err)
		os.Exit(1)
	}

	natsClient.AddClosedCB(func(conn *nats.Conn) {
		logger.Error("Close on NATS client. nats.Conn: ", err, lager.Data{"connection": *conn})
		os.Exit(1)
	})

	return natsClient
}

func InitLoggerFromConfig(logger lager.Logger, c *config.Config, logCounter *vcap.LogCounter) {
	if c.Logging.File != "" {
		file, err := os.OpenFile(c.Logging.File, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
		if err != nil {
			logger.Error("error-opening-file", err, lager.Data{"file": c.Logging.File})
		}
		var logLevel lager.LogLevel
		switch c.Logging.Level {
		case DEBUG:
			logLevel = lager.DEBUG
		case INFO:
			logLevel = lager.INFO
		case ERROR:
			logLevel = lager.ERROR
		case FATAL:
			logLevel = lager.FATAL
		default:
			panic(fmt.Errorf("unknown log level: %s", c.Logging.Level))
		}
		logger.RegisterSink(lager.NewWriterSink(file, logLevel))
	}

	logger.RegisterSink(logCounter)
}
<commit_msg>Updated log message for fetching token<commit_after>package main

import (
	"crypto/tls"

	"github.com/apcera/nats"
	cf_debug_server "github.com/cloudfoundry-incubator/cf-debug-server"
	"github.com/cloudfoundry-incubator/cf-lager"
	routing_api "github.com/cloudfoundry-incubator/routing-api"
	token_fetcher "github.com/cloudfoundry-incubator/uaa-token-fetcher"
	"github.com/cloudfoundry/dropsonde"
	"github.com/cloudfoundry/gorouter/access_log"
	vcap "github.com/cloudfoundry/gorouter/common"
	"github.com/cloudfoundry/gorouter/common/secure"
	"github.com/cloudfoundry/gorouter/config"
	"github.com/cloudfoundry/gorouter/proxy"
	rregistry "github.com/cloudfoundry/gorouter/registry"
	"github.com/cloudfoundry/gorouter/route_fetcher"
	"github.com/cloudfoundry/gorouter/router"
	rvarz "github.com/cloudfoundry/gorouter/varz"
	"github.com/cloudfoundry/yagnats"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"

	"flag"
	"fmt"
	"os"
	"runtime"
	"syscall"
	"time"

	"github.com/cloudfoundry/gorouter/metrics"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"
)

var configFile string

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func main() {
	flag.StringVar(&configFile, "c", "", "Configuration File")
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	logger, _ := cf_lager.New("gorouter.stdout")
	c := config.DefaultConfig(logger)
	logCounter := vcap.NewLogCounter()

	if configFile != "" {
		c = config.InitConfigFromFile(logger, configFile)
	}

	InitLoggerFromConfig(logger, c, logCounter)
	err := dropsonde.Initialize(c.Logging.MetronAddress, c.Logging.JobName)
	if err != nil {
		logger.Error("Dropsonde failed to initialize ", err)
		os.Exit(1)
	}

	// setup number of procs
	if c.GoMaxProcs != 0 {
		runtime.GOMAXPROCS(c.GoMaxProcs)
	}

	if c.DebugAddr != "" {
		cf_debug_server.Run(c.DebugAddr)
	}

	logger.Info("Setting up NATs connection")
	natsClient := connectToNatsServer(logger, c)

	metricsReporter := metrics.NewMetricsReporter()
	registry := rregistry.NewRouteRegistry(c, natsClient, metricsReporter)

	varz := rvarz.NewVarz(registry)
	compositeReporter := metrics.NewCompositeReporter(varz, metricsReporter)

	accessLogger, err := access_log.CreateRunningAccessLogger(logger, c)
	if err != nil {
		logger.Fatal("Error creating access logger: ", err)
	}

	var crypto secure.Crypto
	var cryptoPrev secure.Crypto
	if c.RouteServiceEnabled {
		crypto = createCrypto(logger, c.RouteServiceSecret)
		if c.RouteServiceSecretPrev != "" {
			cryptoPrev = createCrypto(logger, c.RouteServiceSecretPrev)
		}
	}

	proxy := buildProxy(logger, c, registry, accessLogger, compositeReporter, crypto, cryptoPrev)

	router, err := router.NewRouter(logger, c, proxy, natsClient, registry, varz, logCounter, nil)
	if err != nil {
		logger.Error("An error occurred: ", err)
		os.Exit(1)
	}

	members := grouper.Members{
		{"router", router},
	}
	if c.RoutingApiEnabled() {
		logger.Info("Setting up route fetcher")
		routeFetcher := setupRouteFetcher(logger, c, registry)
		members = append(members, grouper.Member{"router-fetcher", routeFetcher})
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1))

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("gorouter.exited-with-failure: ", err)
		os.Exit(1)
	}

	os.Exit(0)
}

func createCrypto(logger lager.Logger, secret string) *secure.AesGCM {
	// generate secure encryption key using key derivation function (pbkdf2)
	secretPbkdf2 := secure.NewPbkdf2([]byte(secret), 16)
	crypto, err := secure.NewAesGCM(secretPbkdf2)
	if err != nil {
		logger.Error("Error creating route service crypto: %s\n", err)
		os.Exit(1)
	}
	return crypto
}

func buildProxy(logger lager.Logger, c *config.Config, registry rregistry.RegistryInterface, accessLogger access_log.AccessLogger, reporter metrics.ProxyReporter, crypto secure.Crypto, cryptoPrev secure.Crypto) proxy.Proxy {
	args := proxy.ProxyArgs{
		Logger:          logger,
		EndpointTimeout: c.EndpointTimeout,
		Ip:              c.Ip,
		TraceKey:        c.TraceKey,
		Registry:        registry,
		Reporter:        reporter,
		AccessLogger:    accessLogger,
		SecureCookies:   c.SecureCookies,
		TLSConfig: &tls.Config{
			CipherSuites:       c.CipherSuites,
			InsecureSkipVerify: c.SSLSkipValidation,
		},
		RouteServiceEnabled: c.RouteServiceEnabled,
		RouteServiceTimeout: c.RouteServiceTimeout,
		Crypto:              crypto,
		CryptoPrev:          cryptoPrev,
		ExtraHeadersToLog:   c.ExtraHeadersToLog,
	}
	return proxy.NewProxy(args)
}

func setupRouteFetcher(logger lager.Logger, c *config.Config, registry rregistry.RegistryInterface) *route_fetcher.RouteFetcher {
	clock := clock.NewClock()

	tokenFetcher := newTokenFetcher(logger, clock, c)
	_, err := tokenFetcher.FetchToken(false)
	if err != nil {
		logger.Error("Unable to fetch token: ", err)
		os.Exit(1)
	}

	routingApiUri := fmt.Sprintf("%s:%d", c.RoutingApi.Uri, c.RoutingApi.Port)
	routingApiClient := routing_api.NewClient(routingApiUri)

	routeFetcher := route_fetcher.NewRouteFetcher(logger, tokenFetcher, registry, c, routingApiClient, 1, clock)
	return routeFetcher
}

func newTokenFetcher(logger lager.Logger, clock clock.Clock, c *config.Config) token_fetcher.TokenFetcher {
	if c.RoutingApi.AuthDisabled {
		logger.Info("using noop token fetcher")
		return token_fetcher.NewNoOpTokenFetcher()
	}
	tokenFetcherConfig := token_fetcher.TokenFetcherConfig{
		MaxNumberOfRetries:   c.TokenFetcherMaxRetries,
		RetryInterval:        c.TokenFetcherRetryInterval,
		ExpirationBufferTime: c.TokenFetcherExpirationBufferTimeInSeconds,
	}

	tokenFetcher, err := token_fetcher.NewTokenFetcher(logger, &c.OAuth, tokenFetcherConfig, clock)
	if err != nil {
		logger.Error("Error creating token fetcher: %s\n", err)
		os.Exit(1)
	}
	logger.Info("Fetching token from UAA")
	return tokenFetcher
}

func connectToNatsServer(logger lager.Logger, c *config.Config) yagnats.NATSConn {
	var natsClient yagnats.NATSConn
	var err error

	natsServers := c.NatsServers()
	attempts := 3
	for attempts > 0 {
		natsClient, err = yagnats.Connect(natsServers)
		if err == nil {
			break
		} else {
			attempts--
			time.Sleep(100 * time.Millisecond)
		}
	}

	if err != nil {
		logger.Error("Error connecting to NATS: %s\n", err)
		os.Exit(1)
	}

	natsClient.AddClosedCB(func(conn *nats.Conn) {
		logger.Error("Close on NATS client. nats.Conn: ", err, lager.Data{"connection": *conn})
		os.Exit(1)
	})

	return natsClient
}

func InitLoggerFromConfig(logger lager.Logger, c *config.Config, logCounter *vcap.LogCounter) {
	if c.Logging.File != "" {
		file, err := os.OpenFile(c.Logging.File, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0666)
		if err != nil {
			logger.Error("error-opening-file", err, lager.Data{"file": c.Logging.File})
		}
		var logLevel lager.LogLevel
		switch c.Logging.Level {
		case DEBUG:
			logLevel = lager.DEBUG
		case INFO:
			logLevel = lager.INFO
		case ERROR:
			logLevel = lager.ERROR
		case FATAL:
			logLevel = lager.FATAL
		default:
			panic(fmt.Errorf("unknown log level: %s", c.Logging.Level))
		}
		logger.RegisterSink(lager.NewWriterSink(file, logLevel))
	}

	logger.RegisterSink(logCounter)
}
<|endoftext|>