<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/dgrijalva/jwt-go"
	"github.com/gorilla/mux"
	"github.com/sokil/go-statsd-client"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"
)

// Version is a current git commit hash and tag
// Injected by compilation flag
var Version = "Unknown"

// BuildNumber is a current commit hash
// Injected by compilation flag
var BuildNumber = "Unknown"

// BuildDate is a date of build
// Injected by compilation flag
var BuildDate = "Unknown"

// HTTP connection params
const defaultHTTPHost = "127.0.0.1"
const defaultHTTPPort = 80

// StatsD connection params
const defaultStatsDHost = "127.0.0.1"
const defaultStatsDPort = 8125

// JWT params
const jwtHeaderName = "X-JWT-Token"

// declare command line options
var httpHost = flag.String("http-host", defaultHTTPHost, "HTTP Host")
var httpPort = flag.Int("http-port", defaultHTTPPort, "HTTP Port")
var statsdHost = flag.String("statsd-host", defaultStatsDHost, "StatsD Host")
var statsdPort = flag.Int("statsd-port", defaultStatsDPort, "StatsD Port")
var metricPrefix = flag.String("metric-prefix", "", "Prefix of metric name")
var tokenSecret = flag.String("jwt-secret", "", "Secret to encrypt JWT")
var verbose = flag.Bool("verbose", false, "Verbose")
var version = flag.Bool("version", false, "Show version")

// statsd client
var statsdClient *statsd.Client

func main() {
	// get flags
	flag.Parse()

	// show version and exit
	if *version == true {
		showVersion()
		os.Exit(0)
	}

	// configure verbosity of logging
	if *verbose == true {
		log.SetOutput(os.Stderr)
	} else {
		log.SetOutput(ioutil.Discard)
	}

	// prepare metric prefix
	if *metricPrefix != "" && (*metricPrefix)[len(*metricPrefix)-1:] != "." {
		*metricPrefix = *metricPrefix + "."
	}

	// create HTTP router
	router := mux.NewRouter().StrictSlash(true)

	// register http request handlers
	router.Handle(
		"/heartbeat",
		validateCORS(http.HandlerFunc(handleHeartbeatRequest)),
	).Methods("GET")

	router.Handle(
		"/count/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleCountRequest))),
	).Methods("POST")

	router.Handle(
		"/gauge/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleGaugeRequest))),
	).Methods("POST")

	router.Handle(
		"/timing/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleTimingRequest))),
	).Methods("POST")

	router.Handle(
		"/set/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleSetRequest))),
	).Methods("POST")

	// Create a new StatsD connection
	statsdClient = statsd.NewClient(*statsdHost, *statsdPort)
	statsdClient.Open()
	defer statsdClient.Close()

	// get server address to bind
	httpAddress := fmt.Sprintf("%s:%d", *httpHost, *httpPort)
	log.Printf("Starting HTTP server %s", httpAddress)

	// create http server
	s := &http.Server{
		Addr:           httpAddress,
		Handler:        router,
		ReadTimeout:    1 * time.Second,
		WriteTimeout:   1 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	// start http server
	err := s.ListenAndServe()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
}

// validate CORS headers
func validateCORS(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")
		if origin != "" {
			w.Header().Add("Access-Control-Allow-Headers", jwtHeaderName+", X-Requested-With, Origin, Accept, Content-Type, Authentication")
			w.Header().Add("Access-Control-Allow-Methods", "GET, POST, HEAD, OPTIONS")
			w.Header().Add("Access-Control-Allow-Origin", origin)
		}
		next.ServeHTTP(w, r)
	})
}

// validate JWT middleware
func validateJWT(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if *tokenSecret == "" {
			next.ServeHTTP(w, r)
		} else {
			// get JWT
			tokenString := r.Header.Get(jwtHeaderName)
			if tokenString == "" {
				http.Error(w, "Token not specified", 401)
				return
			}

			// parse JWT
			_, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
				if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
					return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
				}
				return []byte(*tokenSecret), nil
			})

			if err != nil {
				http.Error(w, "Error parsing token", 403)
				return
			}

			// accept request
			next.ServeHTTP(w, r)
		}

	})
}

// Handle heartbeat request
func handleHeartbeatRequest(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "OK")
}

// Handle StatsD Count request
func handleCountRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get count value
	var value = 1
	valuePostFormValue := r.PostFormValue("value")
	if valuePostFormValue != "" {
		var err error
		value, err = strconv.Atoi(valuePostFormValue)
		if err != nil {
			http.Error(w, "Invalid value specified", 400)
		}
	}

	// get sample rate
	var sampleRate float64 = 1
	sampleRatePostFormValue := r.PostFormValue("sampleRate")
	if sampleRatePostFormValue != "" {
		var err error
		sampleRate, err = strconv.ParseFloat(sampleRatePostFormValue, 32)
		if err != nil {
			http.Error(w, "Invalid sample rate specified", 400)
		}

	}

	// send request
	statsdClient.Count(key, value, float32(sampleRate))
}

// Handle StatsD Gauge request
func handleGaugeRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get gauge shift
	shiftPostFormValue := r.PostFormValue("shift")
	if shiftPostFormValue != "" {
		// get value
		value, err := strconv.Atoi(shiftPostFormValue)
		if err != nil {
			http.Error(w, "Invalid gauge shift specified", 400)
		}
		// send request
		statsdClient.GaugeShift(key, value)
		return
	}

	// get gauge value
	var value = 1
	valuePostFormValue := r.PostFormValue("value")
	if valuePostFormValue != "" {
		// get value
		var err error
		value, err = strconv.Atoi(valuePostFormValue)
		if err != nil {
			http.Error(w, "Invalid gauge value specified", 400)
		}
	}

	// send gauge value request
	statsdClient.Gauge(key, value)

}

// Handle StatsD Timing request
func handleTimingRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get timing
	time, err := strconv.ParseInt(r.PostFormValue("time"), 10, 64)
	if err != nil {
		http.Error(w, "Invalid time specified", 400)
	}

	// get sample rate
	var sampleRate float64 = 1
	sampleRatePostFormValue := r.PostFormValue("sampleRate")
	if sampleRatePostFormValue != "" {
		var err error
		sampleRate, err = strconv.ParseFloat(sampleRatePostFormValue, 32)
		if err != nil {
			http.Error(w, "Invalid sample rate specified", 400)
		}
	}

	// send request
	statsdClient.Timing(key, time, float32(sampleRate))
}

// Handle StatsD Set request
func handleSetRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get set value
	var value = 1
	valuePostFormValue := r.PostFormValue("value")
	if valuePostFormValue != "" {
		var err error
		value, err = strconv.Atoi(valuePostFormValue)
		if err != nil {
			http.Error(w, "Invalid set value specified", 400)
		}
	}

	// send request
	statsdClient.Set(key, value)
}

func showVersion() {
	fmt.Printf(
		"StatsD HTTP Proxy v.%s, build %s from %s\n",
		Version,
		BuildNumber,
		BuildDate,
	)
}
<commit_msg>handle cors preflight<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/dgrijalva/jwt-go"
	"github.com/gorilla/mux"
	"github.com/sokil/go-statsd-client"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"
)

// Version is a current git commit hash and tag
// Injected by compilation flag
var Version = "Unknown"

// BuildNumber is a current commit hash
// Injected by compilation flag
var BuildNumber = "Unknown"

// BuildDate is a date of build
// Injected by compilation flag
var BuildDate = "Unknown"

// HTTP connection params
const defaultHTTPHost = "127.0.0.1"
const defaultHTTPPort = 80

// StatsD connection params
const defaultStatsDHost = "127.0.0.1"
const defaultStatsDPort = 8125

// JWT params
const jwtHeaderName = "X-JWT-Token"

// declare command line options
var httpHost = flag.String("http-host", defaultHTTPHost, "HTTP Host")
var httpPort = flag.Int("http-port", defaultHTTPPort, "HTTP Port")
var statsdHost = flag.String("statsd-host", defaultStatsDHost, "StatsD Host")
var statsdPort = flag.Int("statsd-port", defaultStatsDPort, "StatsD Port")
var metricPrefix = flag.String("metric-prefix", "", "Prefix of metric name")
var tokenSecret = flag.String("jwt-secret", "", "Secret to encrypt JWT")
var verbose = flag.Bool("verbose", false, "Verbose")
var version = flag.Bool("version", false, "Show version")

// statsd client
var statsdClient *statsd.Client

func main() {
	// get flags
	flag.Parse()

	// show version and exit
	if *version == true {
		showVersion()
		os.Exit(0)
	}

	// configure verbosity of logging
	if *verbose == true {
		log.SetOutput(os.Stderr)
	} else {
		log.SetOutput(ioutil.Discard)
	}

	// prepare metric prefix
	if *metricPrefix != "" && (*metricPrefix)[len(*metricPrefix)-1:] != "." {
		*metricPrefix = *metricPrefix + "."
	}

	// create HTTP router
	router := mux.NewRouter().StrictSlash(true)

	// register http request handlers
	router.Handle(
		"/heartbeat",
		validateCORS(http.HandlerFunc(handleHeartbeatRequest)),
	).Methods("GET")

	router.Handle(
		"/count/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleCountRequest))),
	).Methods("POST")

	router.Handle(
		"/gauge/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleGaugeRequest))),
	).Methods("POST")

	router.Handle(
		"/timing/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleTimingRequest))),
	).Methods("POST")

	router.Handle(
		"/set/{key}",
		validateCORS(validateJWT(http.HandlerFunc(handleSetRequest))),
	).Methods("POST")

	router.PathPrefix("/").Methods("OPTIONS").HandlerFunc(handlePreFlightCORSRequest);

	// Create a new StatsD connection
	statsdClient = statsd.NewClient(*statsdHost, *statsdPort)
	statsdClient.Open()
	defer statsdClient.Close()

	// get server address to bind
	httpAddress := fmt.Sprintf("%s:%d", *httpHost, *httpPort)
	log.Printf("Starting HTTP server at %s", httpAddress)

	// create http server
	s := &http.Server{
		Addr:           httpAddress,
		Handler:        router,
		ReadTimeout:    1 * time.Second,
		WriteTimeout:   1 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	// start http server
	err := s.ListenAndServe()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
}

// validate CORS headers
func validateCORS(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")
		if origin != "" {
			w.Header().Add("Access-Control-Allow-Origin", origin)
		}
		next.ServeHTTP(w, r)
	})
}

// validate JWT middleware
func validateJWT(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if *tokenSecret == "" {
			next.ServeHTTP(w, r)
		} else {
			// get JWT
			tokenString := r.Header.Get(jwtHeaderName)
			if tokenString == "" {
				http.Error(w, "Token not specified", 401)
				return
			}

			// parse JWT
			_, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
				if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
					return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
				}
				return []byte(*tokenSecret), nil
			})

			if err != nil {
				http.Error(w, "Error parsing token", 403)
				return
			}

			// accept request
			next.ServeHTTP(w, r)
		}

	})
}

// Handle heartbeat request
func handleHeartbeatRequest(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "OK")
}

// Handle StatsD Count request
func handleCountRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get count value
	var value = 1
	valuePostFormValue := r.PostFormValue("value")
	if valuePostFormValue != "" {
		var err error
		value, err = strconv.Atoi(valuePostFormValue)
		if err != nil {
			http.Error(w, "Invalid value specified", 400)
		}
	}

	// get sample rate
	var sampleRate float64 = 1
	sampleRatePostFormValue := r.PostFormValue("sampleRate")
	if sampleRatePostFormValue != "" {
		var err error
		sampleRate, err = strconv.ParseFloat(sampleRatePostFormValue, 32)
		if err != nil {
			http.Error(w, "Invalid sample rate specified", 400)
		}

	}

	// send request
	statsdClient.Count(key, value, float32(sampleRate))
}

// Handle StatsD Gauge request
func handleGaugeRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get gauge shift
	shiftPostFormValue := r.PostFormValue("shift")
	if shiftPostFormValue != "" {
		// get value
		value, err := strconv.Atoi(shiftPostFormValue)
		if err != nil {
			http.Error(w, "Invalid gauge shift specified", 400)
		}
		// send request
		statsdClient.GaugeShift(key, value)
		return
	}

	// get gauge value
	var value = 1
	valuePostFormValue := r.PostFormValue("value")
	if valuePostFormValue != "" {
		// get value
		var err error
		value, err = strconv.Atoi(valuePostFormValue)
		if err != nil {
			http.Error(w, "Invalid gauge value specified", 400)
		}
	}

	// send gauge value request
	statsdClient.Gauge(key, value)

}

// Handle StatsD Timing request
func handleTimingRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get timing
	time, err := strconv.ParseInt(r.PostFormValue("time"), 10, 64)
	if err != nil {
		http.Error(w, "Invalid time specified", 400)
	}

	// get sample rate
	var sampleRate float64 = 1
	sampleRatePostFormValue := r.PostFormValue("sampleRate")
	if sampleRatePostFormValue != "" {
		var err error
		sampleRate, err = strconv.ParseFloat(sampleRatePostFormValue, 32)
		if err != nil {
			http.Error(w, "Invalid sample rate specified", 400)
		}
	}

	// send request
	statsdClient.Timing(key, time, float32(sampleRate))
}

// Handle StatsD Set request
func handleSetRequest(w http.ResponseWriter, r *http.Request) {
	// get key
	vars := mux.Vars(r)
	key := *metricPrefix + vars["key"]

	// get set value
	var value = 1
	valuePostFormValue := r.PostFormValue("value")
	if valuePostFormValue != "" {
		var err error
		value, err = strconv.Atoi(valuePostFormValue)
		if err != nil {
			http.Error(w, "Invalid set value specified", 400)
		}
	}

	// send request
	statsdClient.Set(key, value)
}

// Handle PreFlight CORS request with OPTIONS method
func handlePreFlightCORSRequest(w http.ResponseWriter, r *http.Request) {
	origin := r.Header.Get("Origin")
	if origin != "" {
		w.Header().Add("Access-Control-Allow-Origin", origin)
		w.Header().Add("Access-Control-Allow-Headers", jwtHeaderName+", X-Requested-With, Origin, Accept, Content-Type, Authentication")
		w.Header().Add("Access-Control-Allow-Methods", "GET, POST, HEAD, OPTIONS")
	}
}

func showVersion() {
	fmt.Printf(
		"StatsD HTTP Proxy v.%s, build %s from %s\n",
		Version,
		BuildNumber,
		BuildDate,
	)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"encoding/json"
	"flag"
	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/golang-lru"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"
	"bytes"
)

const (
	LETSENCRYPT_CREATE_CERTIFICATE_TIMEOUT = time.Minute
	LETSENCRYPT_PRODUCTION_API_URL         = "https://acme-v01.api.letsencrypt.org/directory"
	LETSENCRYPT_STAGING_API_URL            = "https://acme-staging.api.letsencrypt.org/directory"
	PRIVATE_KEY_BITS                       = 2048
	TRY_COUNT                              = 10
	RETRY_SLEEP                            = time.Second * 5
	STATE_FILEMODE                         = 0600
)

var (
	bindTo            = flag.String("bind-to", ":443", "")
	targetPort        = flag.Int("target-port", 80, "")
	targetConnTimeout = flag.Duration("target-conn-timeout", time.Second, "")
	acmeApiUrl        = flag.String("acme-server", LETSENCRYPT_PRODUCTION_API_URL, "")
	acmeTestServer    = flag.Bool("test", false, "Use test lets encrypt server instead of <acme-server>")
	certDir           = flag.String("cert-dir", "certificates", `Directory for save cached certificates. Set cert-dir=- for disable save certs`)
	certMemCount      = flag.Int("in-memory-cnt", 10000, "How many count of certs cache in memory for prevent parse it from file")
	stateFilePath     = flag.String("state-file", "state.json", "Path to save some state data, for example account key")
	proxyMode         = flag.String("proxy-mode", "http", "Proxy-mode after tls handle (http|tcp).")
	httpRealIPHeader  = flag.String("real-ip-header", "Real-IP", "The header will contain original IP of remote connection. It can be few headers, separated by comma.")
	additionalHeadersParam = flag.String("additional-headers", "X-Forwarded-Proto=https", "Additional headers for proxied requests. Several headers separated by comma.")
)

var (
	realIPHeaderNames [][]byte // IP headers, generated by the proxy, included real IP address
	cutHeaders        [][]byte // internal - all headers, that cut from request (upper case).
	additionalHeaders []byte // prepared additional headers
	acmeService       *acmeStruct
)

type stateStruct struct {
	PrivateKey *rsa.PrivateKey
	changed    bool
}

func main() {
	var err error

	flag.Parse()

	// Init
	logrus.SetLevel(logrus.DebugLevel)

	if *proxyMode != "http" && *proxyMode != "tcp" {
		logrus.Panicf("Unknow proxy mode: %v", *proxyMode)
	}
	logrus.Infof("Proxy mode: %v", *proxyMode)

	for _, line := range strings.Split(*httpRealIPHeader, ",") {
		line = strings.TrimSpace(line)
		if line != "" {
			realIPHeaderNames = append(realIPHeaderNames, []byte(line))
			cutHeaders = append(realIPHeaderNames, []byte(strings.ToUpper(line)))
		}
	}
	cutHeaders = append(cutHeaders, []byte("CONNECTION"))

	for _, addHeader := range strings.Split(*additionalHeadersParam, ",") {
		headerParts := strings.SplitN(addHeader, "=", 2)
		if len(headerParts) > 0 {
			cutHeaders = append(cutHeaders, []byte(strings.ToUpper(headerParts[0])))
		}
		buf := &bytes.Buffer{}
		buf.WriteString(headerParts[0])
		buf.WriteByte(':')
		if len(headerParts) == 2 {
			buf.WriteString(headerParts[1])
		}
		buf.WriteString("\r\n")
		additionalHeaders = append(additionalHeaders, buf.Bytes()...)
	}

	localIPs = getLocalIPs()
	acmeService = &acmeStruct{}
	if *certMemCount > 0 {
		logrus.Infof("Create memory cache for '%v' certificates", *certMemCount)
		certMemCache, err = lru.New(*certMemCount)
		if err != nil {
			logrus.Errorf("Can't create memory cache:", err)
			certMemCache = nil
		}
	} else {
		logrus.Info("Memory cache turned off")
	}
	if *certDir == "-" {
		*certDir = ""
	}

	// init service
	var state stateStruct
	stateBytes, err := ioutil.ReadFile(*stateFilePath)
	if err == nil {
		err = json.Unmarshal(stateBytes, &state)
		if err != nil {
			logrus.Errorf("Can't parse state file '%v': %v", *stateFilePath, err)
		}
	} else {
		logrus.Errorf("Can't read state file '%v': %v", *stateFilePath, err)
	}

	acmeService = &acmeStruct{}
	if *acmeTestServer {
		acmeService.serverAddress = LETSENCRYPT_STAGING_API_URL
	} else {
		acmeService.serverAddress = *acmeApiUrl
	}

	if state.PrivateKey == nil {
		logrus.Info("Generate private keys")
		state.PrivateKey, err = rsa.GenerateKey(cryptorand.Reader, PRIVATE_KEY_BITS)
		state.changed = true
		if err != nil {
			logrus.Panic("Can't generate private key")
		}
	} else {
		logrus.Debugf("Skip generate keys - it was read from state")
	}

	saveState(state)

	acmeService.privateKey = state.PrivateKey

	acmeService.Init()
	acmeService.RegisterEnsure(context.TODO())

	// Start listen
	tcpAddr, err := net.ResolveTCPAddr("tcp", *bindTo)
	if err != nil {
		logrus.Panicf("Can't resolve bind-to address '%v': %v", *bindTo, err)
	}
	logrus.Errorf("Start listen: %v", tcpAddr)

	listener, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(err)
	}
	for {
		tcpConn, err := listener.AcceptTCP()
		if err != nil {
			panic(err)
		}
		go handleTcpConnection(tcpConn)
	}
}

func certificateGet(clientHello *tls.ClientHelloInfo) (cert *tls.Certificate, err error) {
	domain := strings.ToLower(clientHello.ServerName)
	if logrus.GetLevel() >= logrus.DebugLevel {
		logrus.Debugf("Required certificate for domain '%v'", domain)
	}

	if strings.HasSuffix(domain, ACME_DOMAIN_SUFFIX) {
		// force generate new certificate, without caching.
		return acmeService.CreateCertificate(domain)
	}

	cert = certificateCacheGet(domain)

	switch {
	case cert != nil && cert.Leaf != nil && cert.Leaf.NotAfter.Before(time.Now()):
		logrus.Warnf("Expired certificate got from cache for domain '%v'", domain)
		// pass to create new certificate.
	case cert != nil:
		return cert, nil
	default:
		// pass
	}

	cert, err = acmeService.CreateCertificate(domain)
	if err == nil {
		certificateCachePut(domain, cert)
	}
	return cert, err
}

func handleTcpConnection(in *net.TCPConn) {
	target, err := getTargetConn(in)
	if err != nil {
		logrus.Errorf("Can't get target IP/port for '%v': %v", target.String(), err)
		return
	}

	// handle ssl
	tlsConfig := tls.Config{
		GetCertificate: certificateGet,
		MinVersion:     tls.VersionSSL30,
	}
	tlsConn := tls.Server(in, &tlsConfig)
	err = tlsConn.Handshake()
	logrus.Debug("tls ciper:", tlsConn.ConnectionState().CipherSuite)
	if err == nil {
		logrus.Debug("Handshake for incoming:", tlsConn.RemoteAddr().String())
	} else {
		logrus.Infof("Error in tls handshake from '%v':%v", tlsConn.RemoteAddr(), err)
	}

	startProxy(target, tlsConn)
}

func getLocalIPs() (res []net.IP) {
	bindAddr, _ := net.ResolveTCPAddr("tcp", *bindTo)
	if bindAddr.IP.IsUnspecified() || len(bindAddr.IP) == 0 {
		addresses, err := net.InterfaceAddrs()
		if err != nil {
			log.Panic("Can't get local ip addresses:", err)
		}
		res = make([]net.IP, 0, len(addresses))
		for _, addr := range addresses {
			logrus.Debug("Local ip:", addr.String())
			ip, _, err := net.ParseCIDR(addr.String())
			if err == nil {
				res = append(res, ip)
			} else {
				logrus.Errorf("Can't parse local ip '%v': %v", addr.String(), err)
			}
		}
	} else {
		res = []net.IP{bindAddr.IP}
	}
	if logrus.GetLevel() >= logrus.InfoLevel {
		ipStrings := make([]string, len(res))
		for i, addr := range res {
			ipStrings[i] = addr.String()
		}
		logrus.Info("Local ip:", ipStrings)
	}
	return res
}

func getTargetConn(in *net.TCPConn) (net.TCPAddr, error) {
	targetAddrP, err := net.ResolveTCPAddr("tcp", in.LocalAddr().String())
	if err != nil {
		logrus.Errorf("Can't resolve local addr '%v': %v", in.LocalAddr().String(), err)
		return net.TCPAddr{}, err
	}
	targetAddrP.Port = *targetPort
	return *targetAddrP, nil
}

func saveState(state stateStruct) {
	if state.changed {
		logrus.Infof("Saving state to '%v'", *stateFilePath)
	} else {
		logrus.Debug("Skip save state becouse it isn't changed")
		return
	}
	stateBytes, err := json.MarshalIndent(&state, "", "    ")
	if err != nil {
		logrus.Errorf("Can't save state to file '%v': %v", *stateFilePath, err)
		return
	}
	err = ioutil.WriteFile(*stateFilePath+".new", stateBytes, STATE_FILEMODE)
	if err != nil {
		logrus.Errorf("Error while write state bytes to file '%v': %v", *stateFilePath+".new", err)
		return
	}
	err = os.Rename(*stateFilePath, *stateFilePath+".old")
	if err != nil {
		logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath, *stateFilePath+".old", err)
	}
	err = os.Rename(*stateFilePath+".new", *stateFilePath)
	if err != nil {
		logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath+".new", *stateFilePath, err)
	}
}
<commit_msg>fix rename error on create new state file.<commit_after>package main

import (
	"context"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"encoding/json"
	"flag"
	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/golang-lru"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"
	"bytes"
)

const (
	LETSENCRYPT_CREATE_CERTIFICATE_TIMEOUT = time.Minute
	LETSENCRYPT_PRODUCTION_API_URL         = "https://acme-v01.api.letsencrypt.org/directory"
	LETSENCRYPT_STAGING_API_URL            = "https://acme-staging.api.letsencrypt.org/directory"
	PRIVATE_KEY_BITS                       = 2048
	TRY_COUNT                              = 10
	RETRY_SLEEP                            = time.Second * 5
	STATE_FILEMODE                         = 0600
)

var (
	bindTo            = flag.String("bind-to", ":443", "")
	targetPort        = flag.Int("target-port", 80, "")
	targetConnTimeout = flag.Duration("target-conn-timeout", time.Second, "")
	acmeApiUrl        = flag.String("acme-server", LETSENCRYPT_PRODUCTION_API_URL, "")
	acmeTestServer    = flag.Bool("test", false, "Use test lets encrypt server instead of <acme-server>")
	certDir           = flag.String("cert-dir", "certificates", `Directory for save cached certificates. Set cert-dir=- for disable save certs`)
	certMemCount      = flag.Int("in-memory-cnt", 10000, "How many count of certs cache in memory for prevent parse it from file")
	stateFilePath     = flag.String("state-file", "state.json", "Path to save some state data, for example account key")
	proxyMode         = flag.String("proxy-mode", "http", "Proxy-mode after tls handle (http|tcp).")
	httpRealIPHeader  = flag.String("real-ip-header", "Real-IP", "The header will contain original IP of remote connection. It can be few headers, separated by comma.")
	additionalHeadersParam = flag.String("additional-headers", "X-Forwarded-Proto=https", "Additional headers for proxied requests. Several headers separated by comma.")
)

var (
	realIPHeaderNames [][]byte // IP headers, generated by the proxy, included real IP address
	cutHeaders        [][]byte // internal - all headers, that cut from request (upper case).
	additionalHeaders []byte // prepared additional headers
	acmeService       *acmeStruct
)

type stateStruct struct {
	PrivateKey *rsa.PrivateKey
	changed    bool
}

func main() {
	var err error

	flag.Parse()

	// Init
	logrus.SetLevel(logrus.DebugLevel)

	if *proxyMode != "http" && *proxyMode != "tcp" {
		logrus.Panicf("Unknow proxy mode: %v", *proxyMode)
	}
	logrus.Infof("Proxy mode: %v", *proxyMode)

	for _, line := range strings.Split(*httpRealIPHeader, ",") {
		line = strings.TrimSpace(line)
		if line != "" {
			realIPHeaderNames = append(realIPHeaderNames, []byte(line))
			cutHeaders = append(realIPHeaderNames, []byte(strings.ToUpper(line)))
		}
	}
	cutHeaders = append(cutHeaders, []byte("CONNECTION"))

	for _, addHeader := range strings.Split(*additionalHeadersParam, ",") {
		headerParts := strings.SplitN(addHeader, "=", 2)
		if len(headerParts) > 0 {
			cutHeaders = append(cutHeaders, []byte(strings.ToUpper(headerParts[0])))
		}
		buf := &bytes.Buffer{}
		buf.WriteString(headerParts[0])
		buf.WriteByte(':')
		if len(headerParts) == 2 {
			buf.WriteString(headerParts[1])
		}
		buf.WriteString("\r\n")
		additionalHeaders = append(additionalHeaders, buf.Bytes()...)
	}

	localIPs = getLocalIPs()
	acmeService = &acmeStruct{}
	if *certMemCount > 0 {
		logrus.Infof("Create memory cache for '%v' certificates", *certMemCount)
		certMemCache, err = lru.New(*certMemCount)
		if err != nil {
			logrus.Errorf("Can't create memory cache:", err)
			certMemCache = nil
		}
	} else {
		logrus.Info("Memory cache turned off")
	}
	if *certDir == "-" {
		*certDir = ""
	}

	// init service
	var state stateStruct
	stateBytes, err := ioutil.ReadFile(*stateFilePath)
	if err == nil {
		err = json.Unmarshal(stateBytes, &state)
		if err != nil {
			logrus.Errorf("Can't parse state file '%v': %v", *stateFilePath, err)
		}
	} else {
		logrus.Errorf("Can't read state file '%v': %v", *stateFilePath, err)
	}

	acmeService = &acmeStruct{}
	if *acmeTestServer {
		acmeService.serverAddress = LETSENCRYPT_STAGING_API_URL
	} else {
		acmeService.serverAddress = *acmeApiUrl
	}

	if state.PrivateKey == nil {
		logrus.Info("Generate private keys")
		state.PrivateKey, err = rsa.GenerateKey(cryptorand.Reader, PRIVATE_KEY_BITS)
		state.changed = true
		if err != nil {
			logrus.Panic("Can't generate private key")
		}
	} else {
		logrus.Debugf("Skip generate keys - it was read from state")
	}

	saveState(state)

	acmeService.privateKey = state.PrivateKey

	acmeService.Init()
	acmeService.RegisterEnsure(context.TODO())

	// Start listen
	tcpAddr, err := net.ResolveTCPAddr("tcp", *bindTo)
	if err != nil {
		logrus.Panicf("Can't resolve bind-to address '%v': %v", *bindTo, err)
	}
	logrus.Errorf("Start listen: %v", tcpAddr)

	listener, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(err)
	}
	for {
		tcpConn, err := listener.AcceptTCP()
		if err != nil {
			panic(err)
		}
		go handleTcpConnection(tcpConn)
	}
}

func certificateGet(clientHello *tls.ClientHelloInfo) (cert *tls.Certificate, err error) {
	domain := strings.ToLower(clientHello.ServerName)
	if logrus.GetLevel() >= logrus.DebugLevel {
		logrus.Debugf("Required certificate for domain '%v'", domain)
	}

	if strings.HasSuffix(domain, ACME_DOMAIN_SUFFIX) {
		// force generate new certificate, without caching.
		return acmeService.CreateCertificate(domain)
	}

	cert = certificateCacheGet(domain)

	switch {
	case cert != nil && cert.Leaf != nil && cert.Leaf.NotAfter.Before(time.Now()):
		logrus.Warnf("Expired certificate got from cache for domain '%v'", domain)
		// pass to create new certificate.
	case cert != nil:
		return cert, nil
	default:
		// pass
	}

	cert, err = acmeService.CreateCertificate(domain)
	if err == nil {
		certificateCachePut(domain, cert)
	}
	return cert, err
}

func handleTcpConnection(in *net.TCPConn) {
	target, err := getTargetConn(in)
	if err != nil {
		logrus.Errorf("Can't get target IP/port for '%v': %v", target.String(), err)
		return
	}

	// handle ssl
	tlsConfig := tls.Config{
		GetCertificate: certificateGet,
		MinVersion:     tls.VersionSSL30,
	}
	tlsConn := tls.Server(in, &tlsConfig)
	err = tlsConn.Handshake()
	logrus.Debug("tls ciper:", tlsConn.ConnectionState().CipherSuite)
	if err == nil {
		logrus.Debug("Handshake for incoming:", tlsConn.RemoteAddr().String())
	} else {
		logrus.Infof("Error in tls handshake from '%v':%v", tlsConn.RemoteAddr(), err)
	}

	startProxy(target, tlsConn)
}

func getLocalIPs() (res []net.IP) {
	bindAddr, _ := net.ResolveTCPAddr("tcp", *bindTo)
	if bindAddr.IP.IsUnspecified() || len(bindAddr.IP) == 0 {
		addresses, err := net.InterfaceAddrs()
		if err != nil {
			log.Panic("Can't get local ip addresses:", err)
		}
		res = make([]net.IP, 0, len(addresses))
		for _, addr := range addresses {
			logrus.Debug("Local ip:", addr.String())
			ip, _, err := net.ParseCIDR(addr.String())
			if err == nil {
				res = append(res, ip)
			} else {
				logrus.Errorf("Can't parse local ip '%v': %v", addr.String(), err)
			}
		}
	} else {
		res = []net.IP{bindAddr.IP}
	}
	if logrus.GetLevel() >= logrus.InfoLevel {
		ipStrings := make([]string, len(res))
		for i, addr := range res {
			ipStrings[i] = addr.String()
		}
		logrus.Info("Local ip:", ipStrings)
	}
	return res
}

func getTargetConn(in *net.TCPConn) (net.TCPAddr, error) {
	targetAddrP, err := net.ResolveTCPAddr("tcp", in.LocalAddr().String())
	if err != nil {
		logrus.Errorf("Can't resolve local addr '%v': %v", in.LocalAddr().String(), err)
		return net.TCPAddr{}, err
	}
	targetAddrP.Port = *targetPort
	return *targetAddrP, nil
}

func saveState(state stateStruct) {
	if state.changed {
		logrus.Infof("Saving state to '%v'", *stateFilePath)
	} else {
		logrus.Debug("Skip save state becouse it isn't changed")
		return
	}
	stateBytes, err := json.MarshalIndent(&state, "", "    ")
	if err != nil {
		logrus.Errorf("Can't save state to file '%v': %v", *stateFilePath, err)
		return
	}
	err = ioutil.WriteFile(*stateFilePath+".new", stateBytes, STATE_FILEMODE)
	if err != nil {
		logrus.Errorf("Error while write state bytes to file '%v': %v", *stateFilePath+".new", err)
		return
	}

	if _, err := os.Stat(*stateFilePath); !os.IsNotExist(err) {
		logrus.Infof("Rename current state file '%v' -> '%v'", *stateFilePath, *stateFilePath + ".old")

		err = os.Rename(*stateFilePath, *stateFilePath + ".old")
		if err != nil {
			logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath, *stateFilePath + ".old", err)
		}
	}

	err = os.Rename(*stateFilePath+".new", *stateFilePath)
	if err != nil {
		logrus.Errorf("Can't rename '%v' to '%v': %v", *stateFilePath+".new", *stateFilePath, err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"

	"github.com/ainoya/sewer/drainer"
	"github.com/ainoya/sewer/flusher"
	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "drain",
			Value: "github",
			Usage: "Destination where you want to notify piped messages",
		},
	}

	var d drainer.Drainer

	app.Action = func(c *cli.Context) error {
		info, _ := os.Stdin.Stat()

		if (info.Mode() & os.ModeCharDevice) == os.ModeCharDevice {
			fmt.Println("The command is intended to work with pipes.")
			fmt.Println("Usage:")
			fmt.Println("cat yourfile.txt | sewer --drain=github")
			return nil
		} else if info.Size() <= 0 {
			fmt.Println("Pipe input is null")
			return nil
		}

		reader := bufio.NewReader(os.Stdin)

		if c.String("drain") == "github" {
			var err error
			d, err = drainer.NewGitHubDrainer()
			if err != nil {
				fmt.Print(err)
				return nil
			}
		} else {
			return nil
		}

		f := flusher.NewFlusher(d, reader)
		err := f.Flush()

		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		return nil
	}

	app.Run(os.Args)
}
<commit_msg>remove input size check because its depends on terminal settings<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"

	"github.com/ainoya/sewer/drainer"
	"github.com/ainoya/sewer/flusher"
	"github.com/urfave/cli"
)

func main() {
	app := cli.NewApp()

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "drain",
			Value: "github",
			Usage: "Destination where you want to notify piped messages",
		},
	}

	var d drainer.Drainer

	app.Action = func(c *cli.Context) error {
		info, _ := os.Stdin.Stat()

		if (info.Mode() & os.ModeCharDevice) == os.ModeCharDevice {
			fmt.Println("The command is intended to work with pipes.")
			fmt.Println("Usage:")
			fmt.Println("cat yourfile.txt | sewer --drain=github")
			return nil
		}

		reader := bufio.NewReader(os.Stdin)

		if c.String("drain") == "github" {
			var err error
			d, err = drainer.NewGitHubDrainer()
			if err != nil {
				fmt.Print(err)
				return nil
			}
		} else {
			return nil
		}

		f := flusher.NewFlusher(d, reader)
		err := f.Flush()

		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		return nil
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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 (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	dockerclient "github.com/docker/docker/client"
	"github.com/pkg/errors"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
)

var (
	podMap = newRegistry() // TODO(ahmetb) globals hurt
)

func main() {
	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		signalChan := make(chan os.Signal, 1)
		signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
		<-signalChan
		cancel()
	}()

	tagEvent := filters.NewArgs()
	tagEvent.Add("type", "image")
	tagEvent.Add("event", "tag")

	k8s, err := kubernetesClient()
	if err != nil {
		log.Fatal(err)
	}
	k8sv, err := k8s.ServerVersion()
	if err != nil {
		log.Fatal("failed to connect to kubernetes")
	}
	log.Printf("connected kubernetes apiserver (%s)", k8sv.GitVersion)

	d, err := dockerclient.NewEnvClient()
	if err != nil {
		log.Fatal(errors.Wrap(err, "cannot create docker client"))
	}
	dv, err := d.ServerVersion(ctx)
	if err != nil {
		log.Fatal(errors.Wrap(err, "failed to connect to docker api"))
	}
	log.Printf("connected docker api (api: v%s, version: %s)", dv.APIVersion, dv.Version)

	podWatcher := podWatchController(k8s)
	go podWatcher.Run(ctx.Done())

	ch, errCh := d.Events(ctx, types.EventsOptions{Filters: tagEvent})
	for {
		select {
		case err := <-errCh:
			select {
			case <-ctx.Done():
				log.Println("Received cancellation.")
				os.Exit(0)
			default:
				panic(err) // TODO(ahmetb) handle gracefully
			}

		case e := <-ch:
			// tag will be in format IMAGE:TAG or IMAGE:latest as it comes
			// from the Docker API (v1.32 at the time of writing).
			tag := e.Actor.Attributes["name"]
			go func() {
				if err := handleTag(k8s, tag); err != nil {
					log.Println(err)
				}
			}()
		case <-ctx.Done():
			break
		}
	}
}

func podWatchController(k8s *kubernetes.Clientset) cache.Controller {
	restClient := k8s.CoreV1().RESTClient()
	lw := cache.NewListWatchFromClient(restClient, "pods", corev1.NamespaceAll, fields.Everything())
	_, controller := cache.NewInformer(lw,
		&corev1.Pod{},
		time.Second*5,
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				pod, ok := obj.(*corev1.Pod)
				if !ok {
					log.Fatalf("list/watch returned non-pod object: %T", obj)
				}
				go trackPod(pod)
			},
			DeleteFunc: func(obj interface{}) {
				pod, ok := obj.(*corev1.Pod)
				if !ok {
					log.Fatalf("list/watch returned non-pod object: %T", obj)
				}
				go untrackPod(pod)
			},
		},
	)
	return controller
}

func handleTag(k8s *kubernetes.Clientset, tag string) error {
	log.Printf("[TAG] %q", tag)
	pods := podMap.get(tag)
	for _, p := range pods {
		log.Printf("[DELETE] pod: %s/%s", p.namespace, p.name)
		if err := k8s.CoreV1().Pods(p.namespace).Delete(p.name, nil); err != nil {
			return errors.Wrap(err, "failed to delete pod")
		}

		// TODO(ahmetb) termination seems to be taking 45 seconds even though
		// terminationGracePeriodSeconds=30. So calling podRegistry here
		// directly to terminate excessive API calls to delete the pod that's
		// already being terminated.
		podMap.del(p, tag)
	}
	return nil
}

func trackPod(p *corev1.Pod) {
	// TODO(ahmetb) implement
	log.Printf("[TRACK] pod %s/%s", p.GetNamespace(), p.GetName())
	for _, c := range p.Spec.Containers {
		podMap.add(pod{
			namespace: p.Namespace,
			name:      p.Name}, canonicalImage(c.Image))
	}
}

func untrackPod(p *corev1.Pod) {
	// TODO(ahmetb) implement
	log.Printf("[UNTRACK] pod %s/%s", p.GetNamespace(), p.GetName())
	for _, c := range p.Spec.Containers {
		podMap.del(pod{
			namespace: p.Namespace,
			name:      p.Name}, canonicalImage(c.Image))
	}
}

// canonicalImage adds :latest to the image tags so
func canonicalImage(img string) string {
	if !strings.Contains(img, ":") {
		// TODO(ahmetb) find better ways to add :latest. currently this detection
		// covers both "IMAGE:TAG" format and "IMAGE@sha256:DIGEST" formats.
		return fmt.Sprintf("%s:latest", img)
	}
	return img
}
<commit_msg>log message cleanup<commit_after>// Copyright 2017 Google Inc.
//
// 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
//
//     https://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 (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	dockerclient "github.com/docker/docker/client"
	"github.com/pkg/errors"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
)

var (
	podMap = newRegistry() // TODO(ahmetb) globals hurt
)

func main() {
	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		signalChan := make(chan os.Signal, 1)
		signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
		sig := <-signalChan
		log.Printf("received signal: %s", sig.String())
		cancel()
	}()

	tagEvent := filters.NewArgs()
	tagEvent.Add("type", "image")
	tagEvent.Add("event", "tag")

	k8s, err := kubernetesClient()
	if err != nil {
		log.Fatal(err)
	}
	k8sv, err := k8s.ServerVersion()
	if err != nil {
		log.Fatal("failed to connect to kubernetes")
	}
	log.Printf("connected kubernetes apiserver (%s)", k8sv.GitVersion)

	d, err := dockerclient.NewEnvClient()
	if err != nil {
		log.Fatal(errors.Wrap(err, "cannot create docker client"))
	}
	dv, err := d.ServerVersion(ctx)
	if err != nil {
		log.Fatal(errors.Wrap(err, "failed to connect to docker api"))
	}
	log.Printf("connected docker api (api: v%s, version: %s)", dv.APIVersion, dv.Version)

	podWatcher := podWatchController(k8s)
	go podWatcher.Run(ctx.Done())

	ch, errCh := d.Events(ctx, types.EventsOptions{Filters: tagEvent})
	for {
		select {
		case err := <-errCh:
			select {
			case <-ctx.Done():
				log.Println("stopping event listener due to cancellation")
				os.Exit(0)
			default:
				panic(err) // TODO(ahmetb) handle gracefully
			}

		case e := <-ch:
			// tag will be in format IMAGE:TAG or IMAGE:latest as it comes
			// from the Docker API (v1.32 at the time of writing).
			tag := e.Actor.Attributes["name"]
			go func() {
				if err := handleTag(k8s, tag); err != nil {
					log.Println(err)
				}
			}()
		case <-ctx.Done():
			break
		}
	}
}

func podWatchController(k8s *kubernetes.Clientset) cache.Controller {
	restClient := k8s.CoreV1().RESTClient()
	lw := cache.NewListWatchFromClient(restClient, "pods", corev1.NamespaceAll, fields.Everything())
	_, controller := cache.NewInformer(lw,
		&corev1.Pod{},
		time.Second*5,
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				pod, ok := obj.(*corev1.Pod)
				if !ok {
					log.Fatalf("list/watch returned non-pod object: %T", obj)
				}
				go trackPod(pod)
			},
			DeleteFunc: func(obj interface{}) {
				pod, ok := obj.(*corev1.Pod)
				if !ok {
					log.Fatalf("list/watch returned non-pod object: %T", obj)
				}
				go untrackPod(pod)
			},
		},
	)
	return controller
}

func handleTag(k8s *kubernetes.Clientset, tag string) error {
	log.Printf("[image_tagged] %q", tag)
	pods := podMap.get(tag)
	for _, p := range pods {
		log.Printf("[delete_pod] %s/%s", p.namespace, p.name)
		if err := k8s.CoreV1().Pods(p.namespace).Delete(p.name, nil); err != nil {
			return errors.Wrap(err, "failed to delete pod")
		}

		// TODO(ahmetb) termination seems to be taking 45 seconds even though
		// terminationGracePeriodSeconds=30. So calling podRegistry here
		// directly to terminate excessive API calls to delete the pod that's
		// already being terminated.
		podMap.del(p, tag)
	}
	return nil
}

func trackPod(p *corev1.Pod) {
	// TODO(ahmetb) implement
	log.Printf("[track_pod] %s/%s", p.GetNamespace(), p.GetName())
	for _, c := range p.Spec.Containers {
		podMap.add(pod{
			namespace: p.Namespace,
			name:      p.Name}, canonicalImage(c.Image))
	}
}

func untrackPod(p *corev1.Pod) {
	log.Printf("[untrack_pod] %s/%s", p.GetNamespace(), p.GetName())
	for _, c := range p.Spec.Containers {
		podMap.del(pod{
			namespace: p.Namespace,
			name:      p.Name}, canonicalImage(c.Image))
	}
}

// canonicalImage adds :latest to the image tags so
func canonicalImage(img string) string {
	if !strings.Contains(img, ":") {
		// TODO(ahmetb) find better ways to add :latest. currently this detection
		// covers both "IMAGE:TAG" format and "IMAGE@sha256:DIGEST" formats.
		return fmt.Sprintf("%s:latest", img)
	}
	return img
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/jawher/mow.cli"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
)

var httpClient = http.Client{
	Transport: &http.Transport{
		MaxIdleConnsPerHost: 128,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
	},
}

func main() {
	app := cli.App("elb dns registrator", "Registers elb cname to *-up.ft.com cnames in dyn using konstructor")

	domains := app.String(cli.StringOpt{
		Name:   "domains",
		Desc:   "comma separated *-up domains",
		EnvVar: "DOMAINS",
	})
	konstructorBaseURL := app.String(cli.StringOpt{
		Name:   "konstructor-base-url",
		Desc:   "konstructor base url: https://dns-api.in.ft.com/v2",
		EnvVar: "KONSTRUCTOR_BASE_URL",
		Value:  "https://dns-api.in.ft.com/v2",
	})
	konstructorAPIKey := app.String(cli.StringOpt{
		Name:   "konstructor-api-key",
		Desc:   "konstructor api key",
		EnvVar: "KONSTRUCTOR_API_KEY",
	})
	elbName := app.String(cli.StringOpt{
		Name:   "elb-name",
		Desc:   "elb cname",
		EnvVar: "ELB_NAME",
	})
	awsRegion := app.String(cli.StringOpt{
		Name:   "aws-region",
		Desc:   "aws region",
		EnvVar: "AWS_REGION",
	})
	awsAccessKeyID := app.String(cli.StringOpt{
		Name:   "aws_access_key_id",
		Desc:   "aws access key id",
		EnvVar: "AWS_ACCESS_KEY_ID",
	})
	awsSecretAccessKey := app.String(cli.StringOpt{
		Name:   "aws_secret_access_key",
		Desc:   "aws secret access key",
		EnvVar: "AWS_SECRET_ACCESS_KEY",
	})

	app.Action = func() {
		c := &conf{
			konsAPIKey:      *konstructorAPIKey,
			konsDNSEndPoint: *konstructorBaseURL,
			elbName:         *elbName,
			awsAccessKey:    *awsAccessKeyID,
			awsSecretKey:    *awsSecretAccessKey,
			awsRegion:       *awsRegion,
		}

		elbCNAME := elbDNSName(c)
		domainsToRegister := strings.Split(*domains, ",")

		for _, domain := range domainsToRegister {
			currentCNAME, err := getCurrentCNAME(c, domain)
			if err != nil {
				log.Fatalf("ERROR - [%v]", err)
			}
			if currentCNAME == "" {
				err = createDNS(c, elbCNAME, domain)
			} else {
				err = updateDNS(c, currentCNAME, elbCNAME, domain)
			}
			if err != nil {
				log.Fatalf("ERROR - [%v]", err)
			}
		}
	}

	err := app.Run(os.Args)
	if err != nil {
		log.Fatalf("ERROR - [%v]", err)
	}
}

type conf struct {
	konsAPIKey      string
	konsDNSEndPoint string
	elbName         string
	awsAccessKey    string
	awsSecretKey    string
	awsRegion       string
}

func elbDNSName(c *conf) string {
	//weirdness around how aws handles credentials
	os.Setenv("AWS_ACCESS_KEY_ID", c.awsAccessKey)
	os.Setenv("AWS_SECRET_ACCESS_KEY", c.awsSecretKey)

	svc := elb.New(
		session.New(
			&aws.Config{
				Region:      aws.String(c.awsRegion),
				Credentials: credentials.NewEnvCredentials(),
			},
		),
	)

	params := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{
			aws.String(c.elbName), // Required
		},
	}

	resp, err := svc.DescribeLoadBalancers(params)

	if err != nil {
		log.Fatal(err)
	}

	return *resp.LoadBalancerDescriptions[0].DNSName
}

func getCurrentCNAME(c *conf, domain string) (string, error) {
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/name/ft.com/%s", c.konsDNSEndPoint, domain), nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept", "application/json")
	req.Header.Add("x-api-key", c.konsAPIKey)

	response, err := httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("Could not connect to Konstructor, [%v]", err)
	}
	defer func() {
		io.Copy(ioutil.Discard, response.Body)
		response.Body.Close()
	}()

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", fmt.Errorf("Could not read konstructor response body, statusCode=[%v], [%v]", response.StatusCode, err)
	}
	if response.StatusCode != http.StatusOK {
		// if status is not 200, log it, it means domain does not exist
		log.Printf("INFO - Domain=[%v] not created, statusCode=[%v], message=[%v]", domain, response.StatusCode, string(data))
		return "", nil
	}

	type konstructorRes struct {
		CNAMES []string `json:"records"`
	}
	r := konstructorRes{}
	if err := json.Unmarshal(data, &r); err != nil {
		return "", err
	}

	//remove trailing "." otherwise update method complains
	return strings.TrimSuffix(r.CNAMES[0], "."), nil
}

func createDNS(c *conf, elbCNAME string, domain string) error {
	body := fmt.Sprintf("{\"zone\": \"ft.com\", \"name\": \"%s\",\"rdata\": \"%s\",\"ttl\": \"30\",\"emailAddress\": \"euan.finlay@ft.com\"}", domain, elbCNAME)
	req, err := http.NewRequest(http.MethodPost, c.konsDNSEndPoint, strings.NewReader(body))
	if err != nil {
		return err
	}
	if err = executeReq(req, c.konsAPIKey); err != nil {
		return fmt.Errorf("Creating domain=[%v] failed, %v", domain, err)
	}
	return nil
}

func updateDNS(c *conf, oldCname string, newCname, domain string) error {
	body := fmt.Sprintf("{\"zone\": \"ft.com\", \"name\": \"%s\",\"oldRdata\": \"%s\",\"newRdata\": \"%s\",\"ttl\": \"30\",\"emailAddress\": \"euan.finlay@ft.com\"}", domain, oldCname, newCname)
	req, err := http.NewRequest(http.MethodPut, c.konsDNSEndPoint, strings.NewReader(body))
	if err != nil {
		return err
	}
	if err = executeReq(req, c.konsAPIKey); err != nil {
		return fmt.Errorf("Updating domain=[%v] failed, %v", domain, err)
	}
	return nil
}

func executeReq(req *http.Request, key string) error {
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.Header.Add("x-api-key", key)

	response, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("Could not connect to Konstructor, [%v]", err)
	}
	defer func() {
		io.Copy(ioutil.Discard, response.Body)
		response.Body.Close()
	}()

	if response.StatusCode != http.StatusOK {
		// if status is not 200, log it, but do not consider it as a service failure
		data, err := ioutil.ReadAll(response.Body)
		message := "Response message could not be obtained"
		if err == nil {
			message = string(data)
		}
		return fmt.Errorf("statusCode=[%v], message=[%v]", response.StatusCode, message)
	}
	return nil
}
<commit_msg>updating DNS deletion email to use UPP group email, not my personal email<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/jawher/mow.cli"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
)

var httpClient = http.Client{
	Transport: &http.Transport{
		MaxIdleConnsPerHost: 128,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
	},
}

func main() {
	app := cli.App("elb dns registrator", "Registers elb cname to *-up.ft.com cnames in dyn using konstructor")

	domains := app.String(cli.StringOpt{
		Name:   "domains",
		Desc:   "comma separated *-up domains",
		EnvVar: "DOMAINS",
	})
	konstructorBaseURL := app.String(cli.StringOpt{
		Name:   "konstructor-base-url",
		Desc:   "konstructor base url: https://dns-api.in.ft.com/v2",
		EnvVar: "KONSTRUCTOR_BASE_URL",
		Value:  "https://dns-api.in.ft.com/v2",
	})
	konstructorAPIKey := app.String(cli.StringOpt{
		Name:   "konstructor-api-key",
		Desc:   "konstructor api key",
		EnvVar: "KONSTRUCTOR_API_KEY",
	})
	elbName := app.String(cli.StringOpt{
		Name:   "elb-name",
		Desc:   "elb cname",
		EnvVar: "ELB_NAME",
	})
	awsRegion := app.String(cli.StringOpt{
		Name:   "aws-region",
		Desc:   "aws region",
		EnvVar: "AWS_REGION",
	})
	awsAccessKeyID := app.String(cli.StringOpt{
		Name:   "aws_access_key_id",
		Desc:   "aws access key id",
		EnvVar: "AWS_ACCESS_KEY_ID",
	})
	awsSecretAccessKey := app.String(cli.StringOpt{
		Name:   "aws_secret_access_key",
		Desc:   "aws secret access key",
		EnvVar: "AWS_SECRET_ACCESS_KEY",
	})

	app.Action = func() {
		c := &conf{
			konsAPIKey:      *konstructorAPIKey,
			konsDNSEndPoint: *konstructorBaseURL,
			elbName:         *elbName,
			awsAccessKey:    *awsAccessKeyID,
			awsSecretKey:    *awsSecretAccessKey,
			awsRegion:       *awsRegion,
		}

		elbCNAME := elbDNSName(c)
		domainsToRegister := strings.Split(*domains, ",")

		for _, domain := range domainsToRegister {
			currentCNAME, err := getCurrentCNAME(c, domain)
			if err != nil {
				log.Fatalf("ERROR - [%v]", err)
			}
			if currentCNAME == "" {
				err = createDNS(c, elbCNAME, domain)
			} else {
				err = updateDNS(c, currentCNAME, elbCNAME, domain)
			}
			if err != nil {
				log.Fatalf("ERROR - [%v]", err)
			}
		}
	}

	err := app.Run(os.Args)
	if err != nil {
		log.Fatalf("ERROR - [%v]", err)
	}
}

type conf struct {
	konsAPIKey      string
	konsDNSEndPoint string
	elbName         string
	awsAccessKey    string
	awsSecretKey    string
	awsRegion       string
}

func elbDNSName(c *conf) string {
	//weirdness around how aws handles credentials
	os.Setenv("AWS_ACCESS_KEY_ID", c.awsAccessKey)
	os.Setenv("AWS_SECRET_ACCESS_KEY", c.awsSecretKey)

	svc := elb.New(
		session.New(
			&aws.Config{
				Region:      aws.String(c.awsRegion),
				Credentials: credentials.NewEnvCredentials(),
			},
		),
	)

	params := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{
			aws.String(c.elbName), // Required
		},
	}

	resp, err := svc.DescribeLoadBalancers(params)

	if err != nil {
		log.Fatal(err)
	}

	return *resp.LoadBalancerDescriptions[0].DNSName
}

func getCurrentCNAME(c *conf, domain string) (string, error) {
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/name/ft.com/%s", c.konsDNSEndPoint, domain), nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept", "application/json")
	req.Header.Add("x-api-key", c.konsAPIKey)

	response, err := httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("Could not connect to Konstructor, [%v]", err)
	}
	defer func() {
		io.Copy(ioutil.Discard, response.Body)
		response.Body.Close()
	}()

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", fmt.Errorf("Could not read konstructor response body, statusCode=[%v], [%v]", response.StatusCode, err)
	}
	if response.StatusCode != http.StatusOK {
		// if status is not 200, log it, it means domain does not exist
		log.Printf("INFO - Domain=[%v] not created, statusCode=[%v], message=[%v]", domain, response.StatusCode, string(data))
		return "", nil
	}

	type konstructorRes struct {
		CNAMES []string `json:"records"`
	}
	r := konstructorRes{}
	if err := json.Unmarshal(data, &r); err != nil {
		return "", err
	}

	//remove trailing "." otherwise update method complains
	return strings.TrimSuffix(r.CNAMES[0], "."), nil
}

func createDNS(c *conf, elbCNAME string, domain string) error {
	body := fmt.Sprintf("{\"zone\": \"ft.com\", \"name\": \"%s\",\"rdata\": \"%s\",\"ttl\": \"30\",\"emailAddress\": \"universal.publishing.platform@ft.com\"}", domain, elbCNAME)
	req, err := http.NewRequest(http.MethodPost, c.konsDNSEndPoint, strings.NewReader(body))
	if err != nil {
		return err
	}
	if err = executeReq(req, c.konsAPIKey); err != nil {
		return fmt.Errorf("Creating domain=[%v] failed, %v", domain, err)
	}
	return nil
}

func updateDNS(c *conf, oldCname string, newCname, domain string) error {
	body := fmt.Sprintf("{\"zone\": \"ft.com\", \"name\": \"%s\",\"oldRdata\": \"%s\",\"newRdata\": \"%s\",\"ttl\": \"30\",\"emailAddress\": \"universal.publishing.platform@ft.com\"}", domain, oldCname, newCname)
	req, err := http.NewRequest(http.MethodPut, c.konsDNSEndPoint, strings.NewReader(body))
	if err != nil {
		return err
	}
	if err = executeReq(req, c.konsAPIKey); err != nil {
		return fmt.Errorf("Updating domain=[%v] failed, %v", domain, err)
	}
	return nil
}

func executeReq(req *http.Request, key string) error {
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.Header.Add("x-api-key", key)

	response, err := httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("Could not connect to Konstructor, [%v]", err)
	}
	defer func() {
		io.Copy(ioutil.Discard, response.Body)
		response.Body.Close()
	}()

	if response.StatusCode != http.StatusOK {
		// if status is not 200, log it, but do not consider it as a service failure
		data, err := ioutil.ReadAll(response.Body)
		message := "Response message could not be obtained"
		if err == nil {
			message = string(data)
		}
		return fmt.Errorf("statusCode=[%v], message=[%v]", response.StatusCode, message)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"database/sql"
	"flag"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"strings"

	"golang.org/x/net/context"

	"github.com/kusubooru/shimmie"
	"github.com/kusubooru/shimmie/store"
)

//go:generate go run generate/version.go

var (
	httpAddr = flag.String("http", "localhost:8080", "HTTP listen address")
	dbDriver = flag.String("dbdriver", "mysql", "database driver")
	dbConfig = flag.String("dbconfig", "", "username:password@(host:port)/database?parseTime=true")
	loginURL = flag.String("loginurl", "/kanri/login", "login URL path to redirect to")
	//version  = flag.Bool("v", false, "print program version")
	certFile = flag.String("tlscert", "", "TLS public key in PEM format.  Must be used together with -tlskey")
	keyFile  = flag.String("tlskey", "", "TLS private key in PEM format.  Must be used together with -tlscert")
	// Set after flag parsing based on certFile & keyFile.
	useTLS bool
)

const (
	description = `Usage: kanri [options]
  Management tools for shimmie.
Options:
`
)

func usage() {
	fmt.Fprintf(os.Stderr, description)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n")
}

func main() {
	flag.Usage = usage
	flag.Parse()
	useTLS = *certFile != "" && *keyFile != ""

	//if *version {
	//	// TheVersion is generated by go generate.
	//	fmt.Printf("teian%v\n", TheVersion)
	//	return
	//}

	// open store with new database connection and create new Shimmie
	shim := shimmie.New(store.Open(*dbDriver, *dbConfig))

	// get common conf
	common, cerr := shim.Store.GetCommon()
	if cerr != nil {
		log.Fatalln("could not get common conf:", cerr)
	}

	// app with Shimmie and common conf
	app := App{Shimmie: shim, Common: common}

	// new context
	ctx := context.Background()

	http.Handle("/kanri", shim.Auth(ctx, app.serveIndex, *loginURL))
	http.Handle("/kanri/safe", shim.Auth(ctx, app.serveSafe, *loginURL))
	http.Handle("/kanri/login", newHandler(ctx, app.serveLogin))
	http.Handle("/kanri/login/submit", newHandler(ctx, app.handleLogin))
	http.Handle("/kanri/logout", http.HandlerFunc(handleLogout))

	if useTLS {
		if err := http.ListenAndServeTLS(*httpAddr, *certFile, *keyFile, nil); err != nil {
			log.Fatalf("Could not start listening (TLS) on %v: %v", *httpAddr, err)
		}
	} else {
		if err := http.ListenAndServe(*httpAddr, nil); err != nil {
			log.Fatalf("Could not start listening on %v: %v", *httpAddr, err)
		}
	}
}

type ctxHandlerFunc func(context.Context, http.ResponseWriter, *http.Request)

func newHandler(ctx context.Context, fn ctxHandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		fn(ctx, w, r)
	}
}

type App struct {
	Shimmie *shimmie.Shimmie
	Common  *shimmie.Common
}

func (app *App) serveIndex(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	app.render(w, indexTmpl, nil)
}

func (app *App) serveSafe(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	images, err := app.Shimmie.GetSafeBustedImages("kusubooru")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	app.render(w, safeTmpl, images)
}

func (app *App) serveLogin(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	app.render(w, loginTmpl, nil)
}

func (app *App) handleLogin(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	// only accept POST method
	if r.Method != "POST" {
		http.Error(w, fmt.Sprintf("%v method not allowed", r.Method), http.StatusMethodNotAllowed)
		return
	}
	username := r.PostFormValue("username")
	password := r.PostFormValue("password")
	user, err := app.Shimmie.GetUser(username)
	if err != nil {
		if err == sql.ErrNoRows {
			app.render(w, loginTmpl, "User does not exist.")
			return
		}
		msg := fmt.Sprintf("get user %q failed: %v", username, err.Error())
		log.Print(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	hash := md5.Sum([]byte(username + password))
	passwordHash := fmt.Sprintf("%x", hash)
	if user.Pass != passwordHash {
		app.render(w, loginTmpl, "Username and password do not match.")
		return
	}
	addr := strings.Split(r.RemoteAddr, ":")[0]
	cookieValue := shimmie.CookieValue(passwordHash, addr)
	shimmie.SetCookie(w, "shm_user", username)
	shimmie.SetCookie(w, "shm_session", cookieValue)
	http.Redirect(w, r, "/kanri", http.StatusFound)
}

func handleLogout(w http.ResponseWriter, r *http.Request) {
	shimmie.SetCookie(w, "shm_user", "")
	shimmie.SetCookie(w, "shm_session", "")
	http.Redirect(w, r, "/kanri", http.StatusFound)
}

func render(w http.ResponseWriter, t *template.Template, data interface{}) {
	if err := t.Execute(w, data); err != nil {
		msg := fmt.Sprintln("could not render template:", err)
		log.Print(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
}

func (app *App) render(w http.ResponseWriter, t *template.Template, data interface{}) {
	render(w, t, struct {
		Data   interface{}
		Common *shimmie.Common
	}{
		Data:   data,
		Common: app.Common,
	})
}
<commit_msg>Add handlers to serve image and thumb<commit_after>package main

import (
	"crypto/md5"
	"database/sql"
	"flag"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"strings"

	"golang.org/x/net/context"

	"github.com/kusubooru/shimmie"
	"github.com/kusubooru/shimmie/store"
)

//go:generate go run generate/version.go

var (
	httpAddr = flag.String("http", "localhost:8080", "HTTP listen address")
	dbDriver = flag.String("dbdriver", "mysql", "database driver")
	dbConfig = flag.String("dbconfig", "", "username:password@(host:port)/database?parseTime=true")
	loginURL = flag.String("loginurl", "/kanri/login", "login URL path to redirect to")
	//version  = flag.Bool("v", false, "print program version")
	certFile = flag.String("tlscert", "", "TLS public key in PEM format.  Must be used together with -tlskey")
	keyFile  = flag.String("tlskey", "", "TLS private key in PEM format.  Must be used together with -tlscert")
	// Set after flag parsing based on certFile & keyFile.
	useTLS bool
)

const (
	description = `Usage: kanri [options]
  Management tools for shimmie.
Options:
`
)

func usage() {
	fmt.Fprintf(os.Stderr, description)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\n")
}

func main() {
	flag.Usage = usage
	flag.Parse()
	useTLS = *certFile != "" && *keyFile != ""

	//if *version {
	//	// TheVersion is generated by go generate.
	//	fmt.Printf("teian%v\n", TheVersion)
	//	return
	//}

	// open store with new database connection and create new Shimmie
	shim := shimmie.New(store.Open(*dbDriver, *dbConfig))

	// get common conf
	common, cerr := shim.Store.GetCommon()
	if cerr != nil {
		log.Fatalln("could not get common conf:", cerr)
	}

	// app with Shimmie and common conf
	app := App{Shimmie: shim, Common: common}

	// new context
	ctx := context.Background()

	http.Handle("/kanri", shim.Auth(ctx, app.serveIndex, *loginURL))
	http.Handle("/kanri/safe", shim.Auth(ctx, app.serveSafe, *loginURL))
	http.Handle("/kanri/_image/", shim.Auth(ctx, app.serveImage, *loginURL))
	http.Handle("/kanri/_thumb/", shim.Auth(ctx, app.serveThumb, *loginURL))
	http.Handle("/kanri/login", newHandler(ctx, app.serveLogin))
	http.Handle("/kanri/login/submit", newHandler(ctx, app.handleLogin))
	http.Handle("/kanri/logout", http.HandlerFunc(handleLogout))

	if useTLS {
		if err := http.ListenAndServeTLS(*httpAddr, *certFile, *keyFile, nil); err != nil {
			log.Fatalf("Could not start listening (TLS) on %v: %v", *httpAddr, err)
		}
	} else {
		if err := http.ListenAndServe(*httpAddr, nil); err != nil {
			log.Fatalf("Could not start listening on %v: %v", *httpAddr, err)
		}
	}
}

type ctxHandlerFunc func(context.Context, http.ResponseWriter, *http.Request)

func newHandler(ctx context.Context, fn ctxHandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		fn(ctx, w, r)
	}
}

type App struct {
	Shimmie *shimmie.Shimmie
	Common  *shimmie.Common
}

func (app *App) serveIndex(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	app.render(w, indexTmpl, nil)
}

func (app *App) serveSafe(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	images, err := app.Shimmie.GetSafeBustedImages("kusubooru")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	app.render(w, safeTmpl, images)
}

func (app *App) serveImage(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	app.serveImageFile(ctx, w, r, app.Shimmie.ImagePath)
}

func (app *App) serveThumb(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	app.serveImageFile(ctx, w, r, app.Shimmie.ThumbPath)
}

func (app *App) serveImageFile(ctx context.Context, w http.ResponseWriter, r *http.Request, path string) {
	hash := r.URL.Path[strings.LastIndex(r.URL.Path, "/")+1:]

	err := app.Shimmie.WriteImageFile(w, path, hash)
	if err != nil {
		http.Error(w, fmt.Sprintf("could not write image file: %v", err), http.StatusInternalServerError)
		return
	}
}

func (app *App) serveLogin(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	app.render(w, loginTmpl, nil)
}

func (app *App) handleLogin(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	// only accept POST method
	if r.Method != "POST" {
		http.Error(w, fmt.Sprintf("%v method not allowed", r.Method), http.StatusMethodNotAllowed)
		return
	}
	username := r.PostFormValue("username")
	password := r.PostFormValue("password")
	user, err := app.Shimmie.GetUser(username)
	if err != nil {
		if err == sql.ErrNoRows {
			app.render(w, loginTmpl, "User does not exist.")
			return
		}
		msg := fmt.Sprintf("get user %q failed: %v", username, err.Error())
		log.Print(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	hash := md5.Sum([]byte(username + password))
	passwordHash := fmt.Sprintf("%x", hash)
	if user.Pass != passwordHash {
		app.render(w, loginTmpl, "Username and password do not match.")
		return
	}
	addr := strings.Split(r.RemoteAddr, ":")[0]
	cookieValue := shimmie.CookieValue(passwordHash, addr)
	shimmie.SetCookie(w, "shm_user", username)
	shimmie.SetCookie(w, "shm_session", cookieValue)
	http.Redirect(w, r, "/kanri", http.StatusFound)
}

func handleLogout(w http.ResponseWriter, r *http.Request) {
	shimmie.SetCookie(w, "shm_user", "")
	shimmie.SetCookie(w, "shm_session", "")
	http.Redirect(w, r, "/kanri", http.StatusFound)
}

func render(w http.ResponseWriter, t *template.Template, data interface{}) {
	if err := t.Execute(w, data); err != nil {
		msg := fmt.Sprintln("could not render template:", err)
		log.Print(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
}

func (app *App) render(w http.ResponseWriter, t *template.Template, data interface{}) {
	render(w, t, struct {
		Data   interface{}
		Common *shimmie.Common
	}{
		Data:   data,
		Common: app.Common,
	})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/mpdroog/beanstalkd" //"github.com/maxid/beanstalkd"
	"gopkg.in/gomail.v1"
	"log"
	"os"
	"github.com/mpdroog/smtpw/config"
	"strings"
	"time"
	"github.com/coreos/go-systemd/daemon"
)

const ERR_WAIT_SEC = 5

var errTimedOut = errors.New("timed out")

var verbose bool
var readonly bool
var hostname string
var L *log.Logger

func proc(m config.Email) error {
	conf, ok := config.C.From[m.From]
	if !ok {
		return errors.New("From does not exist: " + m.From)
	}

	host := conf.Hostname
	if host == "" {
		host = hostname
	}

	msg := gomail.NewMessage()
	msg.SetHeader("Message-ID", fmt.Sprintf("<%s@%s>", RandText(32), host))
	msg.SetHeader("X-Mailer", "smtpw")
	msg.SetHeader("X-Priority", "3")
	if conf.Bounce == nil {
		msg.SetHeader("From", conf.Display+" <"+conf.From+">")
	} else {
		// Set bounce handling
		// From receives bounces
		// Human-clients send to Reply-To
		msg.SetHeader("From", fmt.Sprintf("%s <%s>", conf.Display, *conf.Bounce))
		msg.SetHeader("Reply-To", fmt.Sprintf("%s <%s>", conf.Display, conf.From))
	}
	msg.SetHeader("To", m.To...)
	msg.SetHeader("Bcc", conf.Bcc...)
	msg.SetHeader("Subject", m.Subject)
	msg.SetBody("text/plain", m.Text)
	if len(m.Html) > 0 {
		msg.AddAlternative("text/html", m.Html)
	}

	for name, embed := range m.HtmlEmbed {
		raw, e := base64.StdEncoding.DecodeString(embed)
		if e != nil {
			return errors.New("HtmlEmbed: " + name + " is not base64!")
		}
		if !strings.Contains(m.Html, fmt.Sprintf("cid:"+name)) {
			return errors.New("HtmlEmbed: " + name + " is not used in the HTML!")
		}
		msg.Embed(gomail.CreateFile(name, raw))
	}
	for name, attachment := range m.Attachments {
		raw, e := base64.StdEncoding.DecodeString(attachment)
		if e != nil {
			return errors.New("Attachment: " + name + " is not base64!")
		}
		msg.Embed(gomail.CreateFile(name, raw))
	}

	if readonly {
		L.Printf("From: %s <%s>\n", conf.Display, conf.From)
		L.Printf("To: %v\n", m.To)
		L.Printf("Bcc: %v\n", conf.Bcc)
		L.Printf("Subject: %s\n", m.Subject)
		L.Printf("\ntext/plain\n")
		L.Printf(m.Text)
		L.Printf("\ntext/html\n")
		L.Printf(m.Html)
		L.Printf("\n\n")
		return nil
	}

	mailer := gomail.NewMailer(conf.Host, conf.User, conf.Pass, conf.Port)
	return mailer.Send(msg)
}

func connect() (*beanstalkd.BeanstalkdClient, error) {
	queue, e := beanstalkd.Dial(config.C.Beanstalk)
	if e != nil {
		return nil, e
	}
	// Only listen to email queue.
	queue.Use("email")
	if _, e := queue.Watch("email"); e != nil {
		return nil, e
	}
	queue.Ignore("default")
	return queue, nil
}

func main() {
	var (
		configPath string
		skipOne    bool
	)
	L = log.New(os.Stdout, "", log.LstdFlags)

	flag.BoolVar(&verbose, "v", false, "Verbose-mode")
	flag.BoolVar(&skipOne, "s", false, "Delete e-mail on deverr")
	flag.BoolVar(&readonly, "r", false, "Don't email but flush to stdout")
	flag.StringVar(&configPath, "c", "./config.json", "Path to config.json")
	flag.Parse()

	if e := config.Init(configPath); e != nil {
		panic(e)
	}
	if verbose {
		L.Printf("%+v\n", config.C)
	}
	// TODO: Test config before starting?

	queue, e := connect()
	if e != nil {
		panic(e)
	}
	hostname, e = os.Hostname()
	if e != nil {
		panic(e)
	}

	if verbose {
		L.Printf("SMTPw(%s) email-tube (ignoring default)\n", hostname)
	}
	if readonly {
		L.Printf("!! ReadOnly mode !!\n")
	}

	sent, e := daemon.SdNotify(false, "READY=1")
	if e != nil {
		panic(e)
	}
	if !sent {
		L.Printf("SystemD notify NOT sent\n")
        }

	for {
		job, e := queue.Reserve(15 * 60) //15min timeout
		if e != nil {
			if e.Error() == errTimedOut.Error() {
				// Reserve timeout
				continue
			}

			L.Printf("Beanstalkd err: %s\n", e.Error())
			time.Sleep(time.Second * ERR_WAIT_SEC)
			if strings.HasSuffix(e.Error(), "broken pipe") {
				// Beanstalkd down, reconnect!
				q, e := connect()
				if e != nil {
					L.Printf("Reconnect err: %s\n", e.Error())
				}
				if q != nil {
					queue = q
				}
			}
			continue
		}
		if verbose {
			L.Printf("Parse job %d\n", job.Id)
			L.Printf("JSON:\n%s\n", string(job.Data))
		}
		// Parse
		var m config.Email
		if e := json.Unmarshal(job.Data, &m); e != nil {
			// Broken JSON
			if skipOne {
				L.Printf("WARN: Skip job as JSON is invalid\n")
				queue.Delete(job.Id)
				skipOne = false
				continue
			}

			// Ignore decode trouble
			L.Printf("CRIT: Invalid JSON received (msg=%s)\n", e.Error())
			continue
		}

		if e := proc(m); e != nil {
			// 501 Syntax error in parameters or arguments
			// http://www.greenend.org.uk/rjk/tech/smtpreplies.html
			if strings.HasPrefix(e.Error(), "501 ") {
				L.Printf("WARN: Job buried, invalid email address(es)? (msg=%s)\n", e.Error())
				queue.Bury(job.Id, 1)
				continue
			}
			// TODO: Isolate deverr from senderr
			// Processing trouble?
			L.Printf("WARN: Failed sending, retry in 20sec (msg=%s)\n", e.Error())
			time.Sleep(time.Second * 20)
			continue
		}
		queue.Delete(job.Id)
		if verbose {
			L.Printf("Finished job %d", job.Id)
		}
	}
	queue.Quit()
}
<commit_msg>Bugfix. Add skipone-feature to invalid embeds (most frequent case for breaking)<commit_after>package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/mpdroog/beanstalkd" //"github.com/maxid/beanstalkd"
	"gopkg.in/gomail.v1"
	"log"
	"os"
	"github.com/mpdroog/smtpw/config"
	"strings"
	"time"
	"github.com/coreos/go-systemd/daemon"
)

const ERR_WAIT_SEC = 5

var errTimedOut = errors.New("timed out")

var verbose bool
var readonly bool
var hostname string
var L *log.Logger

func proc(m config.Email, skipOne bool) error {
	conf, ok := config.C.From[m.From]
	if !ok {
		return errors.New("From does not exist: " + m.From)
	}

	host := conf.Hostname
	if host == "" {
		host = hostname
	}

	msg := gomail.NewMessage()
	msg.SetHeader("Message-ID", fmt.Sprintf("<%s@%s>", RandText(32), host))
	msg.SetHeader("X-Mailer", "smtpw")
	msg.SetHeader("X-Priority", "3")
	if conf.Bounce == nil {
		msg.SetHeader("From", conf.Display+" <"+conf.From+">")
	} else {
		// Set bounce handling
		// From receives bounces
		// Human-clients send to Reply-To
		msg.SetHeader("From", fmt.Sprintf("%s <%s>", conf.Display, *conf.Bounce))
		msg.SetHeader("Reply-To", fmt.Sprintf("%s <%s>", conf.Display, conf.From))
	}
	msg.SetHeader("To", m.To...)
	msg.SetHeader("Bcc", conf.Bcc...)
	msg.SetHeader("Subject", m.Subject)
	msg.SetBody("text/plain", m.Text)
	if len(m.Html) > 0 {
		msg.AddAlternative("text/html", m.Html)
	}

	for name, embed := range m.HtmlEmbed {
		raw, e := base64.StdEncoding.DecodeString(embed)
		if e != nil {
			if skipOne {
				L.Printf("WARN: HtmlEmbed: " + name + " is not base64!\n")
				return nil
			}
			return errors.New("HtmlEmbed: " + name + " is not base64!")
		}
		if !strings.Contains(m.Html, fmt.Sprintf("cid:"+name)) {
			if skipOne {
				L.Printf("WARN: HtmlEmbed: " + name + " is not used in the HTML!")
				return nil
			}
			return errors.New("HtmlEmbed: " + name + " is not used in the HTML!")
		}
		msg.Embed(gomail.CreateFile(name, raw))
	}
	for name, attachment := range m.Attachments {
		raw, e := base64.StdEncoding.DecodeString(attachment)
		if e != nil {
			if skipOne {
				L.Printf("WARN: Attachment: " + name + " is not base64!")
				return nil
			}
			return errors.New("Attachment: " + name + " is not base64!")
		}
		msg.Embed(gomail.CreateFile(name, raw))
	}

	if readonly {
		L.Printf("From: %s <%s>\n", conf.Display, conf.From)
		L.Printf("To: %v\n", m.To)
		L.Printf("Bcc: %v\n", conf.Bcc)
		L.Printf("Subject: %s\n", m.Subject)
		L.Printf("\ntext/plain\n")
		L.Printf(m.Text)
		L.Printf("\ntext/html\n")
		L.Printf(m.Html)
		L.Printf("\n\n")
		return nil
	}

	mailer := gomail.NewMailer(conf.Host, conf.User, conf.Pass, conf.Port)
	return mailer.Send(msg)
}

func connect() (*beanstalkd.BeanstalkdClient, error) {
	queue, e := beanstalkd.Dial(config.C.Beanstalk)
	if e != nil {
		return nil, e
	}
	// Only listen to email queue.
	queue.Use("email")
	if _, e := queue.Watch("email"); e != nil {
		return nil, e
	}
	queue.Ignore("default")
	return queue, nil
}

func main() {
	var (
		configPath string
		skipOne    bool
	)
	L = log.New(os.Stdout, "", log.LstdFlags)

	flag.BoolVar(&verbose, "v", false, "Verbose-mode")
	flag.BoolVar(&skipOne, "s", false, "Delete e-mail on deverr")
	flag.BoolVar(&readonly, "r", false, "Don't email but flush to stdout")
	flag.StringVar(&configPath, "c", "./config.json", "Path to config.json")
	flag.Parse()

	if e := config.Init(configPath); e != nil {
		panic(e)
	}
	if verbose {
		L.Printf("%+v\n", config.C)
	}
	// TODO: Test config before starting?

	queue, e := connect()
	if e != nil {
		panic(e)
	}
	hostname, e = os.Hostname()
	if e != nil {
		panic(e)
	}

	if verbose {
		L.Printf("SMTPw(%s) email-tube (ignoring default)\n", hostname)
	}
	if readonly {
		L.Printf("!! ReadOnly mode !!\n")
	}

	sent, e := daemon.SdNotify(false, "READY=1")
	if e != nil {
		panic(e)
	}
	if !sent {
		L.Printf("SystemD notify NOT sent\n")
        }

	for {
		job, e := queue.Reserve(15 * 60) //15min timeout
		if e != nil {
			if e.Error() == errTimedOut.Error() {
				// Reserve timeout
				continue
			}

			L.Printf("Beanstalkd err: %s\n", e.Error())
			time.Sleep(time.Second * ERR_WAIT_SEC)
			if strings.HasSuffix(e.Error(), "broken pipe") {
				// Beanstalkd down, reconnect!
				q, e := connect()
				if e != nil {
					L.Printf("Reconnect err: %s\n", e.Error())
				}
				if q != nil {
					queue = q
				}
			}
			continue
		}
		if verbose {
			L.Printf("Parse job %d\n", job.Id)
			L.Printf("JSON:\n%s\n", string(job.Data))
		}
		// Parse
		var m config.Email
		if e := json.Unmarshal(job.Data, &m); e != nil {
			// Broken JSON
			if skipOne {
				L.Printf("WARN: Skip job as JSON is invalid\n")
				queue.Delete(job.Id)
				skipOne = false
				continue
			}

			// Ignore decode trouble
			L.Printf("CRIT: Invalid JSON received (msg=%s)\n", e.Error())
			continue
		}

		if e := proc(m, skipOne); e != nil {
			// 501 Syntax error in parameters or arguments
			// http://www.greenend.org.uk/rjk/tech/smtpreplies.html
			if strings.HasPrefix(e.Error(), "501 ") {
				L.Printf("WARN: Job buried, invalid email address(es)? (msg=%s)\n", e.Error())
				queue.Bury(job.Id, 1)
				continue
			}
			// TODO: Isolate deverr from senderr
			// Processing trouble?
			L.Printf("WARN: Failed sending, retry in 20sec (msg=%s)\n", e.Error())
			time.Sleep(time.Second * 20)
			continue
		}
		queue.Delete(job.Id)
		if verbose {
			L.Printf("Finished job %d", job.Id)
		}
	}
	queue.Quit()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/codegangsta/negroni"
	"github.com/gorilla/context"
	"github.com/phyber/negroni-gzip/gzip"
)

// ResponseWriter wrapper to catch 404s
type html5mode struct {
	w http.ResponseWriter
	r *http.Request
}

var h5m html5mode
var webRoot string

func redir(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	if r.TLS == nil {
		http.Redirect(w, r, "https://"+r.Host+r.RequestURI, http.StatusMovedPermanently)
	} else {
		next(w, r)
	}
}

func main() {
	var err error
	var port, portTLS int
	var logFile string
	var logger *negroni.Logger

	html5mode := flag.Bool("html5mode", false, "On HTTP 404, serve index.html. Used with AngularJS html5mode.")
	flag.StringVar(&webRoot, "d", "public", "root directory of website")
	flag.StringVar(&logFile, "l", "", "log to a file. Defaults to stdout")
	certfile := flag.String("certFile", "", "SSL certificate filename")
	keyfile := flag.String("keyFile", "", "SSL key filename")
	flag.IntVar(&port, "p", 8080, "HTTP port")
	flag.IntVar(&portTLS, "s", 8081, "HTTPS port")
	forceTLS := flag.Bool("forceTLS", false, "Force HTTPS")
	stripPrefix := flag.String("stripPrefix", "", "path to strip from incoming requests")
	flag.Parse()

	_, err = os.Lstat(webRoot)
	if err != nil {
		log.Fatalf("error opening webroot: %s\n", err)
	} else {
		log.Printf("using webroot '%s'\n", webRoot)
	}

	mux := http.NewServeMux()

	if *stripPrefix != "" {
		log.Printf("stripPrefix '%s'\n", *stripPrefix)
		// Handle either case of trailing slash
		stripPrefixTrim := strings.TrimRight(*stripPrefix, "/")
		mux.Handle(stripPrefixTrim+"/", http.StripPrefix(stripPrefixTrim+"/", http.FileServer(http.Dir(webRoot))))
	} else {
		mux.Handle("/", http.FileServer(http.Dir(webRoot)))

	}

	if logFile != "" {
		f, err := os.OpenFile(logFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			log.Fatalf("error opening logfile '%s', %s\n", logFile, err)
		}
		log.Printf("writing to logfile '%s'\n", logFile)
		logger = &negroni.Logger{log.New(f, "[negroni] ", 0)}
	} else {
		logger = negroni.NewLogger()
	}

	n := negroni.New(
		negroni.NewRecovery(),
		logger,
	)
	if logFile != "" {

	}

	if *forceTLS && *certfile != "" && *keyfile != "" {
		log.Printf("Force TLS enabled\n")
		n.Use(negroni.HandlerFunc(redir))
	}
	n.Use(gzip.Gzip(gzip.DefaultCompression))

	if *html5mode {
		log.Printf("HTML5 mode enabled\n")
		n.Use(negroni.HandlerFunc(html5ModeMiddleware))
	}

	n.UseHandler(mux)

	if *certfile != "" && *keyfile != "" {
		go func() {
			log.Printf("HTTPS listening on port %d\n", portTLS)
			// only support TLS (mitigate against POODLE attack)
			config := &tls.Config{MinVersion: tls.VersionTLS10}
			server := &http.Server{Addr: ":" + strconv.Itoa(portTLS), Handler: n, TLSConfig: config}
			log.Fatal(server.ListenAndServeTLS(*certfile, *keyfile))
		}()
	}

	log.Printf("HTTP listening on port %d\n", port)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(port), n))
}

// This should come before any static file-serving middleware
func html5ModeMiddleware(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	h5m.w = rw
	h5m.r = r
	next(&h5m, r)
	//context.Clear(r)
}

func (sr *html5mode) Header() http.Header { return sr.w.Header() }
func (sr *html5mode) Write(d []byte) (int, error) {
	if _, ok := context.GetOk(sr.r, "html5modeWritten"); ok {
		return 0, nil
	} else {
		return sr.w.Write(d)
	}
}

func (sr *html5mode) WriteHeader(status int) {
	if status == 404 {
		if path.Ext(sr.r.URL.Path) == "" {
			// Server contents of index.html if request isn't for a file
			// Required for Angular.js html5mode
			sr.w.Header().Del("Content-Type")
			http.ServeFile(sr.w, sr.r, webRoot+"/index.html")
			context.Set(sr.r, "html5modeWritten", true)
			return
		}
	}
	sr.w.WriteHeader(status)
}
<commit_msg>Use modern ciphers<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/codegangsta/negroni"
	"github.com/gorilla/context"
	"github.com/phyber/negroni-gzip/gzip"
)

// ResponseWriter wrapper to catch 404s
type html5mode struct {
	w http.ResponseWriter
	r *http.Request
}

var h5m html5mode
var webRoot string

func redir(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	if r.TLS == nil {
		http.Redirect(w, r, "https://"+r.Host+r.RequestURI, http.StatusMovedPermanently)
	} else {
		next(w, r)
	}
}

func main() {
	var err error
	var port, portTLS int
	var logFile string
	var logger *negroni.Logger

	html5mode := flag.Bool("html5mode", false, "On HTTP 404, serve index.html. Used with AngularJS html5mode.")
	flag.StringVar(&webRoot, "d", "public", "root directory of website")
	flag.StringVar(&logFile, "l", "", "log to a file. Defaults to stdout")
	certfile := flag.String("certFile", "", "SSL certificate filename")
	keyfile := flag.String("keyFile", "", "SSL key filename")
	flag.IntVar(&port, "p", 8080, "HTTP port")
	flag.IntVar(&portTLS, "s", 8081, "HTTPS port")
	forceTLS := flag.Bool("forceTLS", false, "Force HTTPS")
	stripPrefix := flag.String("stripPrefix", "", "path to strip from incoming requests")
	flag.Parse()

	_, err = os.Lstat(webRoot)
	if err != nil {
		log.Fatalf("error opening webroot: %s\n", err)
	} else {
		log.Printf("using webroot '%s'\n", webRoot)
	}

	mux := http.NewServeMux()

	if *stripPrefix != "" {
		log.Printf("stripPrefix '%s'\n", *stripPrefix)
		// Handle either case of trailing slash
		stripPrefixTrim := strings.TrimRight(*stripPrefix, "/")
		mux.Handle(stripPrefixTrim+"/", http.StripPrefix(stripPrefixTrim+"/", http.FileServer(http.Dir(webRoot))))
	} else {
		mux.Handle("/", http.FileServer(http.Dir(webRoot)))

	}

	if logFile != "" {
		f, err := os.OpenFile(logFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			log.Fatalf("error opening logfile '%s', %s\n", logFile, err)
		}
		log.Printf("writing to logfile '%s'\n", logFile)
		logger = &negroni.Logger{log.New(f, "[negroni] ", 0)}
	} else {
		logger = negroni.NewLogger()
	}

	n := negroni.New(
		negroni.NewRecovery(),
		logger,
	)
	if logFile != "" {

	}

	if *forceTLS && *certfile != "" && *keyfile != "" {
		log.Printf("Force TLS enabled\n")
		n.Use(negroni.HandlerFunc(redir))
	}
	n.Use(gzip.Gzip(gzip.DefaultCompression))

	if *html5mode {
		log.Printf("HTML5 mode enabled\n")
		n.Use(negroni.HandlerFunc(html5ModeMiddleware))
	}

	n.UseHandler(mux)

	if *certfile != "" && *keyfile != "" {
		go func() {
			log.Printf("HTTPS listening on port %d\n", portTLS)
			config := &tls.Config{}
			// only support TLS (mitigate against POODLE attack)
			config.MinVersion = tls.VersionTLS10
			//Use only modern ciphers
			config.CipherSuites = []uint16{tls.TLS_RSA_WITH_AES_128_CBC_SHA,
				tls.TLS_RSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
				tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
				tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
				tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256}
			server := &http.Server{Addr: ":" + strconv.Itoa(portTLS), Handler: n, TLSConfig: config}
			log.Fatal(server.ListenAndServeTLS(*certfile, *keyfile))
		}()
	}

	log.Printf("HTTP listening on port %d\n", port)
	log.Fatal(http.ListenAndServe(":"+strconv.Itoa(port), n))
}

// This should come before any static file-serving middleware
func html5ModeMiddleware(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	h5m.w = rw
	h5m.r = r
	next(&h5m, r)
	//context.Clear(r)
}

func (sr *html5mode) Header() http.Header { return sr.w.Header() }
func (sr *html5mode) Write(d []byte) (int, error) {
	if _, ok := context.GetOk(sr.r, "html5modeWritten"); ok {
		return 0, nil
	} else {
		return sr.w.Write(d)
	}
}

func (sr *html5mode) WriteHeader(status int) {
	if status == 404 {
		if path.Ext(sr.r.URL.Path) == "" {
			// Server contents of index.html if request isn't for a file
			// Required for Angular.js html5mode
			sr.w.Header().Del("Content-Type")
			http.ServeFile(sr.w, sr.r, webRoot+"/index.html")
			context.Set(sr.r, "html5modeWritten", true)
			return
		}
	}
	sr.w.WriteHeader(status)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"io/ioutil"
)

func main() {
	s, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		panic("Error reading standard input")
	}

	l := lex("stdin", string(s))
	p := parse(l)
//	d := p.parseDomain()
	d := p.parseProblem()
	fmt.Printf("%+v\n", d)
}
<commit_msg>Add a dump flag to main to turn off dumping of the domain/problem.<commit_after>package main

import (
	"fmt"
	"os"
	"io/ioutil"
)

const dump = false

func main() {
	s, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		panic("Error reading standard input")
	}

	l := lex("stdin", string(s))
	p := parse(l)
	d := p.parseDomain()
//	d := p.parseProblem()
	if (dump) {
		fmt.Printf("%+v\n", d)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/go-plugins-helpers/volume"
)

const socketAddress = "/run/docker/plugins/sshfs.sock"

type sshfsVolume struct {
	Password string
	Sshcmd   string
	Port     string

	Mountpoint  string
	connections int
}

type sshfsDriver struct {
	sync.RWMutex

	root      string
	statePath string
	volumes   map[string]*sshfsVolume
}

func newSshfsDriver(root string) (*sshfsDriver, error) {
	logrus.WithField("method", "new driver").Debug(root)

	d := &sshfsDriver{
		root:      filepath.Join(root, "volumes"),
		statePath: filepath.Join(root, "state", "sshfs-state.json"),
		volumes:   map[string]*sshfsVolume{},
	}

	data, err := ioutil.ReadFile(d.statePath)
	if err != nil {
		if os.IsNotExist(err) {
			logrus.WithField("statePath", d.statePath).Debug("no state found")
		} else {
			return nil, err
		}
	} else {
		if err := json.Unmarshal(data, &d.volumes); err != nil {
			return nil, err
		}
	}

	return d, nil
}

func (d *sshfsDriver) saveState() {
	data, err := json.Marshal(d.volumes)
	if err != nil {
		logrus.WithField("statePath", d.statePath).Error(err)
		return
	}

	if err := ioutil.WriteFile(d.statePath, data, 0644); err != nil {
		logrus.WithField("savestate", d.statePath).Error(err)
	}
}

func (d *sshfsDriver) Create(r volume.Request) volume.Response {
	logrus.WithField("method", "create").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()
	v := &sshfsVolume{}

	for key, val := range r.Options {
		switch key {
		case "sshcmd":
			v.Sshcmd = val
		case "password":
			v.Password = val
		case "port":
			v.Port = val
		default:
			return responseError(fmt.Sprintf("unknown option %q", val))
		}
	}

	if v.Sshcmd == "" {
		return responseError("'sshcmd' option required")
	}
	v.Mountpoint = filepath.Join(d.root, fmt.Sprintf("%x", md5.Sum([]byte(v.Sshcmd))))

	d.volumes[r.Name] = v

	d.saveState()

	return volume.Response{}
}

func (d *sshfsDriver) Remove(r volume.Request) volume.Response {
	logrus.WithField("method", "remove").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	if v.connections != 0 {
		return responseError(fmt.Sprintf("volume %s is currently used by a container", r.Name))
	}
	if err := os.RemoveAll(v.Mountpoint); err != nil {
		return responseError(err.Error())
	}
	delete(d.volumes, r.Name)
	d.saveState()
	return volume.Response{}
}

func (d *sshfsDriver) Path(r volume.Request) volume.Response {
	logrus.WithField("method", "path").Debugf("%#v", r)

	d.RLock()
	defer d.RUnlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	return volume.Response{Mountpoint: v.Mountpoint}
}

func (d *sshfsDriver) Mount(r volume.MountRequest) volume.Response {
	logrus.WithField("method", "mount").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	if v.connections == 0 {
		fi, err := os.Lstat(v.Mountpoint)
		if os.IsNotExist(err) {
			if err := os.MkdirAll(v.Mountpoint, 0755); err != nil {
				return responseError(err.Error())
			}
		} else if err != nil {
			return responseError(err.Error())
		}

		if fi != nil && !fi.IsDir() {
			return responseError(fmt.Sprintf("%v already exist and it's not a directory", v.Mountpoint))
		}

		if err := d.mountVolume(v); err != nil {
			return responseError(err.Error())
		}
	}

	v.connections++

	return volume.Response{Mountpoint: v.Mountpoint}
}

func (d *sshfsDriver) Unmount(r volume.UnmountRequest) volume.Response {
	logrus.WithField("method", "unmount").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()
	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	v.connections--

	if v.connections <= 0 {
		if err := d.unmountVolume(v.Mountpoint); err != nil {
			return responseError(err.Error())
		}
		v.connections = 0
	}

	return volume.Response{}
}

func (d *sshfsDriver) Get(r volume.Request) volume.Response {
	logrus.WithField("method", "get").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	return volume.Response{Volume: &volume.Volume{Name: r.Name, Mountpoint: v.Mountpoint}}
}

func (d *sshfsDriver) List(r volume.Request) volume.Response {
	logrus.WithField("method", "list").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	var vols []*volume.Volume
	for name, v := range d.volumes {
		vols = append(vols, &volume.Volume{Name: name, Mountpoint: v.Mountpoint})
	}
	return volume.Response{Volumes: vols}
}

func (d *sshfsDriver) Capabilities(r volume.Request) volume.Response {
	logrus.WithField("method", "capabilities").Debugf("%#v", r)

	return volume.Response{Capabilities: volume.Capability{Scope: "local"}}
}

func (d *sshfsDriver) mountVolume(v *sshfsVolume) error {
	cmd := fmt.Sprintf("sshfs -oStrictHostKeyChecking=no %s %s", v.Sshcmd, v.Mountpoint)
	if v.Port != "" {
		cmd = fmt.Sprintf("%s -p %s", cmd, v.Port)
	}
	if v.Password != "" {
		cmd = fmt.Sprintf("echo %s | %s -o workaround=rename -o password_stdin", v.Password, cmd)
	}
	logrus.Debug(cmd)
	return exec.Command("sh", "-c", cmd).Run()
}

func (d *sshfsDriver) unmountVolume(target string) error {
	cmd := fmt.Sprintf("umount %s", target)
	logrus.Debug(cmd)
	return exec.Command("sh", "-c", cmd).Run()
}

func responseError(err string) volume.Response {
	logrus.Error(err)
	return volume.Response{Err: err}
}

func main() {
	debug := os.Getenv("DEBUG")
	if ok, _ := strconv.ParseBool(debug); ok {
		logrus.SetLevel(logrus.DebugLevel)
	}

	d, err := newSshfsDriver("/mnt")
	if err != nil {
		log.Fatal(err)
	}
	h := volume.NewHandler(d)
	logrus.Infof("listening on %s", socketAddress)
	logrus.Error(h.ServeUnix(socketAddress, 0))
}
<commit_msg>Make sshfs shelling-out more robust<commit_after>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/go-plugins-helpers/volume"
)

const socketAddress = "/run/docker/plugins/sshfs.sock"

type sshfsVolume struct {
	Password string
	Sshcmd   string
	Port     string

	Mountpoint  string
	connections int
}

type sshfsDriver struct {
	sync.RWMutex

	root      string
	statePath string
	volumes   map[string]*sshfsVolume
}

func newSshfsDriver(root string) (*sshfsDriver, error) {
	logrus.WithField("method", "new driver").Debug(root)

	d := &sshfsDriver{
		root:      filepath.Join(root, "volumes"),
		statePath: filepath.Join(root, "state", "sshfs-state.json"),
		volumes:   map[string]*sshfsVolume{},
	}

	data, err := ioutil.ReadFile(d.statePath)
	if err != nil {
		if os.IsNotExist(err) {
			logrus.WithField("statePath", d.statePath).Debug("no state found")
		} else {
			return nil, err
		}
	} else {
		if err := json.Unmarshal(data, &d.volumes); err != nil {
			return nil, err
		}
	}

	return d, nil
}

func (d *sshfsDriver) saveState() {
	data, err := json.Marshal(d.volumes)
	if err != nil {
		logrus.WithField("statePath", d.statePath).Error(err)
		return
	}

	if err := ioutil.WriteFile(d.statePath, data, 0644); err != nil {
		logrus.WithField("savestate", d.statePath).Error(err)
	}
}

func (d *sshfsDriver) Create(r volume.Request) volume.Response {
	logrus.WithField("method", "create").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()
	v := &sshfsVolume{}

	for key, val := range r.Options {
		switch key {
		case "sshcmd":
			v.Sshcmd = val
		case "password":
			v.Password = val
		case "port":
			v.Port = val
		default:
			return responseError(fmt.Sprintf("unknown option %q", val))
		}
	}

	if v.Sshcmd == "" {
		return responseError("'sshcmd' option required")
	}
	v.Mountpoint = filepath.Join(d.root, fmt.Sprintf("%x", md5.Sum([]byte(v.Sshcmd))))

	d.volumes[r.Name] = v

	d.saveState()

	return volume.Response{}
}

func (d *sshfsDriver) Remove(r volume.Request) volume.Response {
	logrus.WithField("method", "remove").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	if v.connections != 0 {
		return responseError(fmt.Sprintf("volume %s is currently used by a container", r.Name))
	}
	if err := os.RemoveAll(v.Mountpoint); err != nil {
		return responseError(err.Error())
	}
	delete(d.volumes, r.Name)
	d.saveState()
	return volume.Response{}
}

func (d *sshfsDriver) Path(r volume.Request) volume.Response {
	logrus.WithField("method", "path").Debugf("%#v", r)

	d.RLock()
	defer d.RUnlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	return volume.Response{Mountpoint: v.Mountpoint}
}

func (d *sshfsDriver) Mount(r volume.MountRequest) volume.Response {
	logrus.WithField("method", "mount").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	if v.connections == 0 {
		fi, err := os.Lstat(v.Mountpoint)
		if os.IsNotExist(err) {
			if err := os.MkdirAll(v.Mountpoint, 0755); err != nil {
				return responseError(err.Error())
			}
		} else if err != nil {
			return responseError(err.Error())
		}

		if fi != nil && !fi.IsDir() {
			return responseError(fmt.Sprintf("%v already exist and it's not a directory", v.Mountpoint))
		}

		if err := d.mountVolume(v); err != nil {
			return responseError(err.Error())
		}
	}

	v.connections++

	return volume.Response{Mountpoint: v.Mountpoint}
}

func (d *sshfsDriver) Unmount(r volume.UnmountRequest) volume.Response {
	logrus.WithField("method", "unmount").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()
	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	v.connections--

	if v.connections <= 0 {
		if err := d.unmountVolume(v.Mountpoint); err != nil {
			return responseError(err.Error())
		}
		v.connections = 0
	}

	return volume.Response{}
}

func (d *sshfsDriver) Get(r volume.Request) volume.Response {
	logrus.WithField("method", "get").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	v, ok := d.volumes[r.Name]
	if !ok {
		return responseError(fmt.Sprintf("volume %s not found", r.Name))
	}

	return volume.Response{Volume: &volume.Volume{Name: r.Name, Mountpoint: v.Mountpoint}}
}

func (d *sshfsDriver) List(r volume.Request) volume.Response {
	logrus.WithField("method", "list").Debugf("%#v", r)

	d.Lock()
	defer d.Unlock()

	var vols []*volume.Volume
	for name, v := range d.volumes {
		vols = append(vols, &volume.Volume{Name: name, Mountpoint: v.Mountpoint})
	}
	return volume.Response{Volumes: vols}
}

func (d *sshfsDriver) Capabilities(r volume.Request) volume.Response {
	logrus.WithField("method", "capabilities").Debugf("%#v", r)

	return volume.Response{Capabilities: volume.Capability{Scope: "local"}}
}

func (d *sshfsDriver) mountVolume(v *sshfsVolume) error {
	cmd := exec.Command("sshfs", "-oStrictHostKeyChecking=no", v.Sshcmd, v.Mountpoint)
	if v.Port != "" {
		cmd.Args = append(cmd.Args, "-p", v.Port)
	}
	if v.Password != "" {
		cmd.Args = append(cmd.Args, "-p", v.Port, "-o", "workaround=rename", "-o", "password_stdin")
		cmd.Stdin = strings.NewReader(v.Password)
	}
	logrus.Debug(cmd.Args)
	return cmd.Run()
}

func (d *sshfsDriver) unmountVolume(target string) error {
	cmd := fmt.Sprintf("umount %s", target)
	logrus.Debug(cmd)
	return exec.Command("sh", "-c", cmd).Run()
}

func responseError(err string) volume.Response {
	logrus.Error(err)
	return volume.Response{Err: err}
}

func main() {
	debug := os.Getenv("DEBUG")
	if ok, _ := strconv.ParseBool(debug); ok {
		logrus.SetLevel(logrus.DebugLevel)
	}

	d, err := newSshfsDriver("/mnt")
	if err != nil {
		log.Fatal(err)
	}
	h := volume.NewHandler(d)
	logrus.Infof("listening on %s", socketAddress)
	logrus.Error(h.ServeUnix(socketAddress, 0))
}
<|endoftext|>