<commit_before>package urknall

import (
	"fmt"
	"strings"
)

// The host type. Use the "NewHost" function to create the basic value.
//
// Please note that you need to set the primary interface (the one the host is accessible on) name, if that is not
// "eth0". That should only be necessary on rare circumstances.
//
// A host is added a set of packages, that are provisioned on request.
//
//	TODO(gfrey): Add better support for interfaces and IPs.
type Host struct {
	IP       string // Host's IP address used to provision the system.
	User     string // User used to log in.
	Password string // SSH password to be used (besides ssh-agent)
	Port     int    // SSH Port to be used

	Tags []string
	Env  []string // custom env settings to be used for all sessions

	packageNames []string
	runlists     []*Runlist
}

// Get the user used to access the host. If none is given the "root" account is as default.
func (h *Host) user() string {
	if h.User == "" {
		return "root"
	}
	return h.User
}

// Alias for the AddCommands methods.
func (h *Host) Add(name string, cmd interface{}, cmds ...interface{}) {
	h.AddCommands(name, cmd, cmds...)
}

// Register the list of given commands (either of the cmd.Command type or as string) as a package (without
// configuration) with the given name.
func (h *Host) AddCommands(name string, cmd interface{}, cmds ...interface{}) {
	cmdList := append([]interface{}{cmd}, cmds...)
	h.AddPackage(name, NewPackage(cmdList...))
}

// Add the given package with the given name to the host.
//
// The name is used as reference during provisioning and allows for provisioning the very same package in different
// configuration (with different version for example). Package names must be unique and the "uk." prefix is reserved for
// urknall internal packages.
func (h *Host) AddPackage(name string, pkg Package) {
	if strings.HasPrefix(name, "uk.") {
		panic(fmt.Sprintf(`package name prefix "uk." reserved (in %q)`, name))
	}

	if strings.Contains(name, " ") {
		panic(fmt.Sprintf(`package names must not contain spaces (%q does)`, name))
	}

	for i := range h.packageNames {
		if h.packageNames[i] == name {
			panic(fmt.Sprintf("package with name %q exists already", name))
		}
	}

	h.packageNames = append(h.packageNames, name)
	h.runlists = append(h.runlists, &Runlist{name: name, pkg: pkg, host: h})
}

// Provision the host, i.e. execute all the commands contained in the packages registered with this host.
func (h *Host) Provision(opts *ProvisionOptions) (e error) {
	sc := newSSHClient(h, opts)
	return sc.provision()
}

// Predicate to test whether sudo is required (user for the host is not "root").
func (h *Host) isSudoRequired() bool {
	if h.User != "" && h.User != "root" {
		return true
	}
	return false
}

func (h *Host) precompileRunlists() (e error) {
	for _, runlist := range h.runlists {
		if len(runlist.commands) > 0 {
			return fmt.Errorf("pkg %q seems to be packaged already", runlist.name)
		}

		if e = runlist.compile(); e != nil {
			return e
		}
	}

	return nil
}
<commit_msg>unified variadic arguments usage<commit_after>package urknall

import (
	"fmt"
	"strings"
)

// The host type. Use the "NewHost" function to create the basic value.
//
// Please note that you need to set the primary interface (the one the host is accessible on) name, if that is not
// "eth0". That should only be necessary on rare circumstances.
//
// A host is added a set of packages, that are provisioned on request.
//
//	TODO(gfrey): Add better support for interfaces and IPs.
type Host struct {
	IP       string // Host's IP address used to provision the system.
	User     string // User used to log in.
	Password string // SSH password to be used (besides ssh-agent)
	Port     int    // SSH Port to be used

	Tags []string
	Env  []string // custom env settings to be used for all sessions

	packageNames []string
	runlists     []*Runlist
}

// Get the user used to access the host. If none is given the "root" account is as default.
func (h *Host) user() string {
	if h.User == "" {
		return "root"
	}
	return h.User
}

// Alias for the AddCommands methods.
func (h *Host) Add(name string, cmds ...interface{}) {
	h.AddCommands(name, cmds...)
}

// Register the list of given commands (either of the cmd.Command type or as string) as a package (without
// configuration) with the given name.
func (h *Host) AddCommands(name string, cmds ...interface{}) {
	h.AddPackage(name, NewPackage(cmds...))
}

// Add the given package with the given name to the host.
//
// The name is used as reference during provisioning and allows for provisioning the very same package in different
// configuration (with different version for example). Package names must be unique and the "uk." prefix is reserved for
// urknall internal packages.
func (h *Host) AddPackage(name string, pkg Package) {
	if strings.HasPrefix(name, "uk.") {
		panic(fmt.Sprintf(`package name prefix "uk." reserved (in %q)`, name))
	}

	if strings.Contains(name, " ") {
		panic(fmt.Sprintf(`package names must not contain spaces (%q does)`, name))
	}

	for i := range h.packageNames {
		if h.packageNames[i] == name {
			panic(fmt.Sprintf("package with name %q exists already", name))
		}
	}

	h.packageNames = append(h.packageNames, name)
	h.runlists = append(h.runlists, &Runlist{name: name, pkg: pkg, host: h})
}

// Provision the host, i.e. execute all the commands contained in the packages registered with this host.
func (h *Host) Provision(opts *ProvisionOptions) (e error) {
	sc := newSSHClient(h, opts)
	return sc.provision()
}

// Predicate to test whether sudo is required (user for the host is not "root").
func (h *Host) isSudoRequired() bool {
	if h.User != "" && h.User != "root" {
		return true
	}
	return false
}

func (h *Host) precompileRunlists() (e error) {
	for _, runlist := range h.runlists {
		if len(runlist.commands) > 0 {
			return fmt.Errorf("pkg %q seems to be packaged already", runlist.name)
		}

		if e = runlist.compile(); e != nil {
			return e
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package citadel

import (
	"fmt"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/citadel/citadel/utils"
	"github.com/samalba/dockerclient"
)

// Host represents a host machine on the cluster
// running docker containers
type Host struct {
	mux sync.Mutex

	// ID is a unique identifier for a host
	ID string `json:"id,omitempty"`
	// Cpus is the number of cpus the host has available
	Cpus int `json:"cpus,omitempty"`
	// Memory is the amount of memory in mb the host has available
	Memory int `json:"memory,omitempty"`
	// Label is specific attributes of a host
	Labels []string `json:"labels,omitempty"`

	// containers that were started with citadel
	managedContainers map[string]*Container

	logger *logrus.Logger
	docker *dockerclient.DockerClient
}

func NewHost(id string, cpus, memory int, labels []string, docker *dockerclient.DockerClient, logger *logrus.Logger) (*Host, error) {
	h := &Host{
		ID:                id,
		Cpus:              cpus,
		Memory:            memory,
		Labels:            labels,
		docker:            docker,
		logger:            logger,
		managedContainers: make(map[string]*Container),
	}

	docker.StartMonitorEvents(h.eventHandler, nil)

	return h, nil
}

func (h *Host) eventHandler(event *dockerclient.Event, _ ...interface{}) {
	switch event.Status {
	case "die":
		container, err := h.inspect(event.Id)
		if err != nil {
			h.logger.WithField("error", err).Error("fetch dead container information")
			return
		}
		h.mux.Lock()

		// only restart it if it's a managed container
		if c, exists := h.managedContainers[container.ID]; exists && c.Type == Service {
			container.State.ExitedAt = time.Now()

			if err := h.startContainer(container); err != nil {
				h.logger.WithField("error", err).Error("restarting dead container")
			}
		}

		h.mux.Unlock()
	default:
		h.logger.WithFields(logrus.Fields{
			"type": event.Status,
			"id":   event.Id,
			"from": event.From,
		}).Debug("docker event")
	}
}

// Close stops the events monitor
func (h *Host) Close() error {
	h.mux.Lock()

	h.docker.StopAllMonitorEvents()

	h.mux.Unlock()

	return nil
}

func (h *Host) Containers() []*Container {
	out := []*Container{}
	h.mux.Lock()

	for _, c := range h.managedContainers {
		out = append(out, c)
	}

	h.mux.Unlock()

	return out
}

func (h *Host) RunContainer(c *Container) error {
	h.mux.Lock()
	defer h.mux.Unlock()

	config := &dockerclient.ContainerConfig{
		Image:  c.Image,
		Memory: c.Memory * 1024 * 1024,
		Cpuset: utils.IToCpuset(c.Cpus),
	}

	if _, err := h.docker.CreateContainer(config, c.ID); err != nil {
		return err
	}

	if err := h.startContainer(c); err != nil {
		return err
	}

	return nil
}

func (h *Host) startContainer(c *Container) error {
	var hostConfig *dockerclient.HostConfig

	if c.Ports != nil {
		hostConfig = &dockerclient.HostConfig{
			PortBindings: make(map[string][]dockerclient.PortBinding),
		}

		for _, p := range c.Ports {
			proto := "tcp"
			if p.Proto != "" {
				proto = p.Proto
			}

			hostConfig.PortBindings[fmt.Sprintf("%d/%s", p.Container, proto)] = []dockerclient.PortBinding{
				{
					HostPort: fmt.Sprint(p.Host),
				},
			}
		}
	}

	if err := h.docker.StartContainer(c.ID, hostConfig); err != nil {
		return err
	}

	current, err := h.inspect(c.ID)
	if err != nil {
		return err
	}

	c.State = current.State
	c.State.StartedAt = time.Now()
	c.State.ExitedAt = time.Time{}

	h.managedContainers[c.ID] = c

	return nil
}

func (h *Host) StopContainer(c *Container) error {
	h.mux.Lock()
	defer h.mux.Unlock()

	if err := h.docker.StopContainer(c.ID, 10); err != nil {
		return err
	}

	// update the state on the original container so that when it is
	// returned it has the latest information
	current, err := h.inspect(c.ID)
	if err != nil {
		return err
	}

	c.State = current.State
	c.State.ExitedAt = time.Now()

	delete(h.managedContainers, c.ID)

	return h.docker.RemoveContainer(c.ID)
}

func (h *Host) inspect(id string) (*Container, error) {
	info, err := h.docker.InspectContainer(id)
	if err != nil {
		return nil, err
	}

	return containerFromDocker(h, info)
}
<commit_msg>Don't run already existing containers<commit_after>package citadel

import (
	"fmt"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/citadel/citadel/utils"
	"github.com/samalba/dockerclient"
)

// Host represents a host machine on the cluster
// running docker containers
type Host struct {
	mux sync.Mutex

	// ID is a unique identifier for a host
	ID string `json:"id,omitempty"`
	// Cpus is the number of cpus the host has available
	Cpus int `json:"cpus,omitempty"`
	// Memory is the amount of memory in mb the host has available
	Memory int `json:"memory,omitempty"`
	// Label is specific attributes of a host
	Labels []string `json:"labels,omitempty"`

	// containers that were started with citadel
	managedContainers map[string]*Container

	logger *logrus.Logger
	docker *dockerclient.DockerClient
}

func NewHost(id string, cpus, memory int, labels []string, docker *dockerclient.DockerClient, logger *logrus.Logger) (*Host, error) {
	h := &Host{
		ID:                id,
		Cpus:              cpus,
		Memory:            memory,
		Labels:            labels,
		docker:            docker,
		logger:            logger,
		managedContainers: make(map[string]*Container),
	}

	docker.StartMonitorEvents(h.eventHandler, nil)

	return h, nil
}

func (h *Host) eventHandler(event *dockerclient.Event, _ ...interface{}) {
	switch event.Status {
	case "die":
		container, err := h.inspect(event.Id)
		if err != nil {
			h.logger.WithField("error", err).Error("fetch dead container information")
			return
		}
		h.mux.Lock()

		// only restart it if it's a managed container
		if c, exists := h.managedContainers[container.ID]; exists && c.Type == Service {
			container.State.ExitedAt = time.Now()

			if err := h.startContainer(container); err != nil {
				h.logger.WithField("error", err).Error("restarting dead container")
			}
		}

		h.mux.Unlock()
	default:
		h.logger.WithFields(logrus.Fields{
			"type": event.Status,
			"id":   event.Id,
			"from": event.From,
		}).Debug("docker event")
	}
}

// Close stops the events monitor
func (h *Host) Close() error {
	h.mux.Lock()

	h.docker.StopAllMonitorEvents()

	h.mux.Unlock()

	return nil
}

func (h *Host) Containers() []*Container {
	out := []*Container{}
	h.mux.Lock()

	for _, c := range h.managedContainers {
		out = append(out, c)
	}

	h.mux.Unlock()

	return out
}

func (h *Host) RunContainer(c *Container) error {
	h.mux.Lock()
	defer h.mux.Unlock()

	if _, exists := h.managedContainers[c.ID]; exists {
		return fmt.Errorf("container %s is already managed", c.ID)
	}

	config := &dockerclient.ContainerConfig{
		Image:  c.Image,
		Memory: c.Memory * 1024 * 1024,
		Cpuset: utils.IToCpuset(c.Cpus),
	}

	if _, err := h.docker.CreateContainer(config, c.ID); err != nil {
		return err
	}

	if err := h.startContainer(c); err != nil {
		return err
	}

	h.managedContainers[c.ID] = c

	return nil
}

func (h *Host) startContainer(c *Container) error {
	var hostConfig *dockerclient.HostConfig

	if c.Ports != nil {
		hostConfig = &dockerclient.HostConfig{
			PortBindings: make(map[string][]dockerclient.PortBinding),
		}

		for _, p := range c.Ports {
			proto := "tcp"
			if p.Proto != "" {
				proto = p.Proto
			}

			hostConfig.PortBindings[fmt.Sprintf("%d/%s", p.Container, proto)] = []dockerclient.PortBinding{
				{
					HostPort: fmt.Sprint(p.Host),
				},
			}
		}
	}

	if err := h.docker.StartContainer(c.ID, hostConfig); err != nil {
		return err
	}

	current, err := h.inspect(c.ID)
	if err != nil {
		return err
	}

	c.State = current.State
	c.State.StartedAt = time.Now()
	c.State.ExitedAt = time.Time{}

	return nil
}

func (h *Host) StopContainer(c *Container) error {
	h.mux.Lock()
	defer h.mux.Unlock()

	if err := h.docker.StopContainer(c.ID, 10); err != nil {
		return err
	}

	// update the state on the original container so that when it is
	// returned it has the latest information
	current, err := h.inspect(c.ID)
	if err != nil {
		return err
	}

	c.State = current.State
	c.State.ExitedAt = time.Now()

	delete(h.managedContainers, c.ID)

	return h.docker.RemoveContainer(c.ID)
}

func (h *Host) inspect(id string) (*Container, error) {
	info, err := h.docker.InspectContainer(id)
	if err != nil {
		return nil, err
	}

	return containerFromDocker(h, info)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/litl/galaxy/log"
)

var (
	httpRouter *HostRouter
)

type RequestLogger struct{}

// This works along with the ServiceRegistry, and the individual Services to
// route http requests based on the Host header. The Resgistry hold the mapping
// of VHost names to individual services, and each service has it's own
// ReeverseProxy to fulfill the request.
// HostRouter contains the ReverseProxy http Listener, and has an http.Handler
// to service the requets.
type HostRouter struct {
	sync.Mutex
	// the http frontend
	server *http.Server

	// track our listener so we can kill the server
	listener net.Listener
}

func NewHostRouter() *HostRouter {
	return &HostRouter{}
}

func (r *HostRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	reqId := genId()
	req.Header.Set("X-Request-Id", reqId)
	defer func(start time.Time) {
		log.Printf("id=%s total_duration=%s", reqId, time.Since(start))
	}(time.Now())

	var err error
	host := req.Host
	if strings.Contains(host, ":") {
		host, _, err = net.SplitHostPort(req.Host)
		if err != nil {
			log.Warnf("%s", err)
		}
	}

	svc := Registry.GetVHostService(host)

	if svc != nil && svc.httpProxy != nil {
		// The vhost has a service registered, give it to the proxy
		svc.ServeHTTP(w, req)
		return
	}

	r.adminHandler(w, req)
}

func (r *HostRouter) adminHandler(w http.ResponseWriter, req *http.Request) {
	r.Lock()
	defer r.Unlock()

	if Registry.VHostsLen() == 0 {
		http.Error(w, "no backends available", http.StatusServiceUnavailable)
		return
	}

	// TODO: better status lines
	stats := Registry.Stats()
	for _, svc := range stats {
		if len(svc.VirtualHosts) == 0 {
			continue
		}
		fmt.Fprintf(w, "%v\n", svc.VirtualHosts)
		for _, b := range svc.Backends {
			js, _ := json.Marshal(b)
			fmt.Fprintf(w, "\t%s\n", string(js))
		}
	}

	fmt.Fprintf(w, "\n")
	return
}

// TODO: collect more stats?

// Start the HTTP Router frontend.
// Takes a channel to notify when the listener is started
// to safely synchronize tests.
func (r *HostRouter) Start(ready chan bool) {
	//FIXME: poor locking strategy
	r.Lock()

	log.Printf("HTTP server listening at %s", listenAddr)

	// Proxy acts as http handler:
	// These timeouts for for overall request duration. They don't effect
	// keepalive, but will close an overly slow request.
	r.server = &http.Server{
		Addr:           listenAddr,
		Handler:        r,
		ReadTimeout:    10 * time.Minute,
		WriteTimeout:   10 * time.Minute,
		MaxHeaderBytes: 1 << 20,
	}

	var err error

	// These timeouts are for each individual Read/Write operation
	// These will close keepalive connections too.
	// TODO: configure timeout somewhere
	r.listener, err = newTimeoutListener(listenAddr, 120*time.Second)
	if err != nil {
		log.Errorf("%s", err)
		r.Unlock()
		return
	}

	r.Unlock()
	if ready != nil {
		close(ready)
	}

	// This will log a closed connection error every time we Stop
	// but that's mostly a testing issue.
	log.Errorf("%s", r.server.Serve(r.listener))
}

func (r *HostRouter) Stop() {
	r.listener.Close()
}

func startHTTPServer() {
	//FIXME: this global wg?
	defer wg.Done()
	httpRouter = NewHostRouter()
	httpRouter.Start(nil)
}

func sslRedirect(pr *ProxyRequest) bool {
	if sslOnly && pr.Request.Header.Get("X-Forwarded-Proto") != "https" {
		//TODO: verify RequestURI
		redirLoc := "https://" + pr.Request.Host + pr.Request.RequestURI
		http.Redirect(pr.ResponseWriter, pr.Request, redirLoc, http.StatusMovedPermanently)
		return false
	}

	return true
}

type ErrorPage struct {
	// The Mutex protects access to the body slice, and headers
	// Everything else should be static once the ErrorPage is created.
	sync.Mutex

	Location    string
	StatusCodes []int

	// body contains the cached error page
	body []byte
	// important headers
	header http.Header
}

func (e *ErrorPage) Body() []byte {
	e.Lock()
	defer e.Unlock()
	return e.body
}

func (e *ErrorPage) SetBody(b []byte) {
	e.Lock()
	defer e.Unlock()
	e.body = b
}

func (e *ErrorPage) Header() http.Header {
	e.Lock()
	defer e.Unlock()
	return e.header
}

func (e *ErrorPage) SetHeader(h http.Header) {
	e.Lock()
	defer e.Unlock()
	e.header = h
}

// List of headers we want to cache for ErrorPages
var ErrorHeaders = []string{
	"Content-Type",
	"Content-Encoding",
	"Cache-Control",
	"Last-Modified",
	"Retry-After",
	"Set-Cookie",
}

// ErrorResponse provides a ReverProxy callback to process a response and
// insert custom error pages for a virtual host.
type ErrorResponse struct {
	sync.Mutex

	// map them by status for responses
	pages map[int]*ErrorPage

	// keep this handy to refresh the pages
	client *http.Client
}

func NewErrorResponse(pages map[string][]int) *ErrorResponse {
	errors := &ErrorResponse{
		pages: make(map[int]*ErrorPage),
	}

	// aggressively timeout connections
	errors.client = &http.Client{
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 2 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 2 * time.Second,
		},
		Timeout: 5 * time.Second,
	}

	if pages != nil {
		errors.Update(pages)
	}
	return errors
}

// Get the ErrorPage, returning nil if the page was incomplete.
// We permanently cache error pages and headers once we've seen them.
func (e *ErrorResponse) Get(code int) *ErrorPage {
	e.Lock()
	page, ok := e.pages[code]
	e.Unlock()

	if !ok {
		// this is a code we don't handle
		return nil
	}

	body := page.Body()
	if body != nil {
		return page
	}

	// we haven't successfully fetched this error
	e.fetch(page)
	return page
}

func (e *ErrorResponse) fetch(page *ErrorPage) {
	log.Debugf("Fetching error page from %s", page.Location)
	resp, err := e.client.Get(page.Location)
	if err != nil {
		log.Warnf("Could not fetch %s: %s", page.Location, err.Error())
		return
	}
	defer resp.Body.Close()

	// If the StatusCode matches any of our registered codes, it's OK
	for _, code := range page.StatusCodes {
		if resp.StatusCode == code {
			resp.StatusCode = http.StatusOK
			break
		}
	}

	if resp.StatusCode != http.StatusOK {
		log.Warnf("Server returned %d when fetching %s", resp.StatusCode, page.Location)
		return
	}

	header := make(map[string][]string)
	for _, key := range ErrorHeaders {
		if hdr, ok := resp.Header[key]; ok {
			header[key] = hdr
		}
	}
	// set the headers along with the body below

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Warnf("Error reading response from %s: %s", page.Location, err.Error())
		return
	}

	if len(body) > 0 {
		page.SetHeader(header)
		page.SetBody(body)
		return
	}
	log.Warnf("Empty response from %s", page.Location)
}

// This replaces all existing ErrorPages
func (e *ErrorResponse) Update(pages map[string][]int) {
	e.Lock()
	defer e.Unlock()

	e.pages = make(map[int]*ErrorPage)

	for loc, codes := range pages {
		page := &ErrorPage{
			StatusCodes: codes,
			Location:    loc,
		}

		for _, code := range codes {
			e.pages[code] = page
		}
		go e.fetch(page)
	}
}

func (e *ErrorResponse) CheckResponse(pr *ProxyRequest) bool {

	errPage := e.Get(pr.Response.StatusCode)
	if errPage != nil {
		// load the cached headers
		header := pr.ResponseWriter.Header()
		for key, val := range errPage.Header() {
			header[key] = val
		}

		pr.ResponseWriter.WriteHeader(pr.Response.StatusCode)
		pr.ResponseWriter.Write(errPage.Body())
		return false
	}

	return true
}

func logProxyRequest(pr *ProxyRequest) bool {
	// TODO: we may to be able to switch this off
	if pr == nil {
		return true
	}

	var id, method, clientIP, url, backend, agent string
	var status int

	duration := pr.FinishTime.Sub(pr.StartTime)

	if pr.Request != nil {
		id = pr.Request.Header.Get("X-Request-Id")
		method = pr.Request.Method
		clientIP = pr.Request.RemoteAddr
		url = pr.Request.Host + pr.Request.RequestURI
		agent = pr.Request.UserAgent()
		status = pr.Response.StatusCode
	}

	if pr.Response != nil && pr.Response.Request != nil && pr.Response.Request.URL != nil {
		backend = pr.Response.Request.URL.Host
	}

	err := fmt.Sprintf("%v", pr.ProxyError)

	fmtStr := "id=%s method=%s clientIp=%s url=%s backend=%s status=%d duration=%s agent=%s, err=%s"

	log.Printf(fmtStr, id, method, clientIP, url, backend, status, duration, agent, err)
	return true
}
<commit_msg>Log X-Forwarded-For as client IP in shuttle<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/litl/galaxy/log"
)

var (
	httpRouter *HostRouter
)

type RequestLogger struct{}

// This works along with the ServiceRegistry, and the individual Services to
// route http requests based on the Host header. The Resgistry hold the mapping
// of VHost names to individual services, and each service has it's own
// ReeverseProxy to fulfill the request.
// HostRouter contains the ReverseProxy http Listener, and has an http.Handler
// to service the requets.
type HostRouter struct {
	sync.Mutex
	// the http frontend
	server *http.Server

	// track our listener so we can kill the server
	listener net.Listener
}

func NewHostRouter() *HostRouter {
	return &HostRouter{}
}

func (r *HostRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	reqId := genId()
	req.Header.Set("X-Request-Id", reqId)
	defer func(start time.Time) {
		log.Printf("id=%s total_duration=%s", reqId, time.Since(start))
	}(time.Now())

	var err error
	host := req.Host
	if strings.Contains(host, ":") {
		host, _, err = net.SplitHostPort(req.Host)
		if err != nil {
			log.Warnf("%s", err)
		}
	}

	svc := Registry.GetVHostService(host)

	if svc != nil && svc.httpProxy != nil {
		// The vhost has a service registered, give it to the proxy
		svc.ServeHTTP(w, req)
		return
	}

	r.adminHandler(w, req)
}

func (r *HostRouter) adminHandler(w http.ResponseWriter, req *http.Request) {
	r.Lock()
	defer r.Unlock()

	if Registry.VHostsLen() == 0 {
		http.Error(w, "no backends available", http.StatusServiceUnavailable)
		return
	}

	// TODO: better status lines
	stats := Registry.Stats()
	for _, svc := range stats {
		if len(svc.VirtualHosts) == 0 {
			continue
		}
		fmt.Fprintf(w, "%v\n", svc.VirtualHosts)
		for _, b := range svc.Backends {
			js, _ := json.Marshal(b)
			fmt.Fprintf(w, "\t%s\n", string(js))
		}
	}

	fmt.Fprintf(w, "\n")
	return
}

// TODO: collect more stats?

// Start the HTTP Router frontend.
// Takes a channel to notify when the listener is started
// to safely synchronize tests.
func (r *HostRouter) Start(ready chan bool) {
	//FIXME: poor locking strategy
	r.Lock()

	log.Printf("HTTP server listening at %s", listenAddr)

	// Proxy acts as http handler:
	// These timeouts for for overall request duration. They don't effect
	// keepalive, but will close an overly slow request.
	r.server = &http.Server{
		Addr:           listenAddr,
		Handler:        r,
		ReadTimeout:    10 * time.Minute,
		WriteTimeout:   10 * time.Minute,
		MaxHeaderBytes: 1 << 20,
	}

	var err error

	// These timeouts are for each individual Read/Write operation
	// These will close keepalive connections too.
	// TODO: configure timeout somewhere
	r.listener, err = newTimeoutListener(listenAddr, 120*time.Second)
	if err != nil {
		log.Errorf("%s", err)
		r.Unlock()
		return
	}

	r.Unlock()
	if ready != nil {
		close(ready)
	}

	// This will log a closed connection error every time we Stop
	// but that's mostly a testing issue.
	log.Errorf("%s", r.server.Serve(r.listener))
}

func (r *HostRouter) Stop() {
	r.listener.Close()
}

func startHTTPServer() {
	//FIXME: this global wg?
	defer wg.Done()
	httpRouter = NewHostRouter()
	httpRouter.Start(nil)
}

func sslRedirect(pr *ProxyRequest) bool {
	if sslOnly && pr.Request.Header.Get("X-Forwarded-Proto") != "https" {
		//TODO: verify RequestURI
		redirLoc := "https://" + pr.Request.Host + pr.Request.RequestURI
		http.Redirect(pr.ResponseWriter, pr.Request, redirLoc, http.StatusMovedPermanently)
		return false
	}

	return true
}

type ErrorPage struct {
	// The Mutex protects access to the body slice, and headers
	// Everything else should be static once the ErrorPage is created.
	sync.Mutex

	Location    string
	StatusCodes []int

	// body contains the cached error page
	body []byte
	// important headers
	header http.Header
}

func (e *ErrorPage) Body() []byte {
	e.Lock()
	defer e.Unlock()
	return e.body
}

func (e *ErrorPage) SetBody(b []byte) {
	e.Lock()
	defer e.Unlock()
	e.body = b
}

func (e *ErrorPage) Header() http.Header {
	e.Lock()
	defer e.Unlock()
	return e.header
}

func (e *ErrorPage) SetHeader(h http.Header) {
	e.Lock()
	defer e.Unlock()
	e.header = h
}

// List of headers we want to cache for ErrorPages
var ErrorHeaders = []string{
	"Content-Type",
	"Content-Encoding",
	"Cache-Control",
	"Last-Modified",
	"Retry-After",
	"Set-Cookie",
}

// ErrorResponse provides a ReverProxy callback to process a response and
// insert custom error pages for a virtual host.
type ErrorResponse struct {
	sync.Mutex

	// map them by status for responses
	pages map[int]*ErrorPage

	// keep this handy to refresh the pages
	client *http.Client
}

func NewErrorResponse(pages map[string][]int) *ErrorResponse {
	errors := &ErrorResponse{
		pages: make(map[int]*ErrorPage),
	}

	// aggressively timeout connections
	errors.client = &http.Client{
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 2 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 2 * time.Second,
		},
		Timeout: 5 * time.Second,
	}

	if pages != nil {
		errors.Update(pages)
	}
	return errors
}

// Get the ErrorPage, returning nil if the page was incomplete.
// We permanently cache error pages and headers once we've seen them.
func (e *ErrorResponse) Get(code int) *ErrorPage {
	e.Lock()
	page, ok := e.pages[code]
	e.Unlock()

	if !ok {
		// this is a code we don't handle
		return nil
	}

	body := page.Body()
	if body != nil {
		return page
	}

	// we haven't successfully fetched this error
	e.fetch(page)
	return page
}

func (e *ErrorResponse) fetch(page *ErrorPage) {
	log.Debugf("Fetching error page from %s", page.Location)
	resp, err := e.client.Get(page.Location)
	if err != nil {
		log.Warnf("Could not fetch %s: %s", page.Location, err.Error())
		return
	}
	defer resp.Body.Close()

	// If the StatusCode matches any of our registered codes, it's OK
	for _, code := range page.StatusCodes {
		if resp.StatusCode == code {
			resp.StatusCode = http.StatusOK
			break
		}
	}

	if resp.StatusCode != http.StatusOK {
		log.Warnf("Server returned %d when fetching %s", resp.StatusCode, page.Location)
		return
	}

	header := make(map[string][]string)
	for _, key := range ErrorHeaders {
		if hdr, ok := resp.Header[key]; ok {
			header[key] = hdr
		}
	}
	// set the headers along with the body below

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Warnf("Error reading response from %s: %s", page.Location, err.Error())
		return
	}

	if len(body) > 0 {
		page.SetHeader(header)
		page.SetBody(body)
		return
	}
	log.Warnf("Empty response from %s", page.Location)
}

// This replaces all existing ErrorPages
func (e *ErrorResponse) Update(pages map[string][]int) {
	e.Lock()
	defer e.Unlock()

	e.pages = make(map[int]*ErrorPage)

	for loc, codes := range pages {
		page := &ErrorPage{
			StatusCodes: codes,
			Location:    loc,
		}

		for _, code := range codes {
			e.pages[code] = page
		}
		go e.fetch(page)
	}
}

func (e *ErrorResponse) CheckResponse(pr *ProxyRequest) bool {

	errPage := e.Get(pr.Response.StatusCode)
	if errPage != nil {
		// load the cached headers
		header := pr.ResponseWriter.Header()
		for key, val := range errPage.Header() {
			header[key] = val
		}

		pr.ResponseWriter.WriteHeader(pr.Response.StatusCode)
		pr.ResponseWriter.Write(errPage.Body())
		return false
	}

	return true
}

func logProxyRequest(pr *ProxyRequest) bool {
	// TODO: we may to be able to switch this off
	if pr == nil || pr.Request == nil {
		return true
	}

	var id, method, clientIP, url, backend, agent string
	var status int

	duration := pr.FinishTime.Sub(pr.StartTime)

	id = pr.Request.Header.Get("X-Request-Id")
	method = pr.Request.Method
	url = pr.Request.Host + pr.Request.RequestURI
	agent = pr.Request.UserAgent()
	status = pr.Response.StatusCode

	clientIP = pr.Request.Header.Get("X-Forwarded-For")
	if clientIP == "" {
		clientIP = pr.Request.RemoteAddr
	}

	if pr.Response != nil && pr.Response.Request != nil && pr.Response.Request.URL != nil {
		backend = pr.Response.Request.URL.Host
	}

	err := fmt.Sprintf("%v", pr.ProxyError)

	fmtStr := "id=%s method=%s clientIp=%s url=%s backend=%s status=%d duration=%s agent=%s, err=%s"

	log.Printf(fmtStr, id, method, clientIP, url, backend, status, duration, agent, err)
	return true
}
<|endoftext|>


<commit_before>package lily

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	netUrl "net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

func HTTPStatusCodeIsOk(statusCode int) bool {
	return statusCode > 199 && statusCode < 300
}

func HTTPSetContentTypeJSON(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
}

func HTTPSetContentTypeHTML(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/html; charset=UTF-8")
}

func HTTPRespondStr(w http.ResponseWriter, code int, body string) {
	if len(body) == 0 {
		panic("body must be not empty")
	}
	w.WriteHeader(code)
	fmt.Fprint(w, body)
}

func HTTPRespondJSONObj(w http.ResponseWriter, code int, obj interface{}) {
	HTTPSetContentTypeJSON(w)
	w.WriteHeader(code)
	ErrPanic(json.NewEncoder(w).Encode(obj))
}

func HTTPRespondJSONParseError(w http.ResponseWriter) {
	HTTPRespond400(w, "bad_json", "Fail to parse JSON")
}

func HTTPSendRequest(withJar bool, method, url string, urlParams map[string]string,
	data []byte, timeout time.Duration, headers ...string) (*http.Response, error) {
	var err error
	var req *http.Request
	var jar http.CookieJar

	if data != nil {
		req, err = http.NewRequest(method, url, bytes.NewBuffer(data))
		ErrPanic(err)
	} else {
		req, err = http.NewRequest(method, url, nil)
		ErrPanic(err)
	}

	if urlParams != nil {
		q := netUrl.Values{}
		for k, v := range urlParams {
			q.Add(k, v)
		}
		req.URL.RawQuery = q.Encode()
	}

	for i := 0; (i + 1) < len(headers); i += 2 {
		req.Header.Set(headers[i], headers[i+1])
	}

	if withJar {
		jar, err = cookiejar.New(nil)
		ErrPanic(err)
	}

	client := http.Client{
		Timeout: timeout,
		Jar:     jar,
	}

	return client.Do(req)
}

func HTTPSendRequestReceiveBytes(withJar, errSCode bool, method, url string, urlParams map[string]string,
	data []byte, timeout time.Duration, headers ...string) (int, []byte, error) {
	var res []byte

	resp, err := HTTPSendRequest(withJar, method, url, urlParams, data, timeout, headers...)
	if err != nil {
		return 0, nil, err
	}
	defer resp.Body.Close()

	res, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0, nil, err
	}

	if !HTTPStatusCodeIsOk(resp.StatusCode) {
		if errSCode {
			return resp.StatusCode, res, errors.New(fmt.Sprintf("bad_http_status_code - %d\nbody: %s", resp.StatusCode, string(res)))
		}
		return resp.StatusCode, res, nil
	}

	return resp.StatusCode, res, nil
}

func HTTPSendRequestReceiveString(withJar, errSCode bool, method, url string, urlParams map[string]string,
	data []byte, timeout time.Duration, headers ...string) (int, string, error) {
	sCode, resBytes, err := HTTPSendRequestReceiveBytes(withJar, errSCode, method, url, urlParams, data, timeout, headers...)

	return sCode, string(resBytes), err
}

func HTTPSendRequestReceiveJSONObj(withJar, errSCode bool, method, url string, urlParams map[string]string,
	data []byte, rObj interface{}, timeout time.Duration, headers ...string) (int, []byte, error) {
	sCode, rBytes, err := HTTPSendRequestReceiveBytes(
		withJar, errSCode, method, url, urlParams, data, timeout, headers...)
	if err != nil || !HTTPStatusCodeIsOk(sCode) {
		return sCode, rBytes, err
	}

	err = json.Unmarshal(rBytes, rObj)
	if err != nil {
		return sCode, rBytes, errors.New(fmt.Sprintf("fail_to_parse_json - %s\nbody: %s", err.Error(), string(rBytes)))
	}

	return sCode, rBytes, nil
}

func HTTPRetrieveRequestHostURL(r *http.Request) string {
	scheme := r.Header.Get("X-Forwarded-Proto")
	if scheme == "" {
		if r.TLS == nil {
			scheme = "http"
		} else {
			scheme = "https"
		}
	}
	return scheme + "://" + r.Host
}

func HTTPRetrieveRemoteIP(r *http.Request) (result string) {
	result = ""
	if parts := strings.Split(r.RemoteAddr, ":"); len(parts) == 2 {
		result = parts[0]
	}
	// If we have a forwarded-for header, take the address from there
	if xff := strings.Trim(r.Header.Get("X-Forwarded-For"), ","); len(xff) > 0 {
		addrs := strings.Split(xff, ",")
		lastFwd := addrs[len(addrs)-1]
		if ip := net.ParseIP(lastFwd); ip != nil {
			result = ip.String()
		}
		// parse X-Real-Ip header
	} else if xri := r.Header.Get("X-Real-Ip"); len(xri) > 0 {
		if ip := net.ParseIP(xri); ip != nil {
			result = ip.String()
		}
	}
	return
}

func HTTPUploadFileFromRequestForm(r *http.Request, key, dirPath, dir string, filename string) (string, error) {
	var err error

	finalDirPath := filepath.Join(dirPath, dir)

	err = os.MkdirAll(finalDirPath, os.ModePerm)
	if err != nil {
		return "", err
	}

	srcFile, header, err := r.FormFile(key)
	if err != nil {
		return "", err
	}
	defer srcFile.Close()

	fileExt := filepath.Ext(header.Filename)
	if fileExt == "" {
		return "", errors.New("bad_extension")
	}

	dstFile, err := TempFile(finalDirPath, filename+"_*"+fileExt)
	if err != nil {
		return "", err
	}
	defer dstFile.Close()

	_, err = io.Copy(dstFile, srcFile)
	if err != nil {
		return "", err
	}

	err = os.Chmod(dstFile.Name(), 0644)
	if err != nil {
		return "", err
	}

	newName, err := filepath.Rel(dirPath, dstFile.Name())
	if err != nil {
		return "", err
	}

	return newName, nil
}

func HTTPRespondError(w http.ResponseWriter, code int, err string, detail string, extras ...interface{}) {
	obj := map[string]interface{}{}
	obj["error"] = err
	obj["error_dsc"] = detail
	for i := 0; (i + 1) < len(extras); i += 2 {
		obj[extras[i].(string)] = extras[i+1]
	}
	HTTPRespondJSONObj(w, code, obj)
}

func HTTPRespond400(w http.ResponseWriter, err, detail string, extras ...interface{}) {
	HTTPRespondError(w, 400, err, detail, extras...)
}

func HTTPRespond401(w http.ResponseWriter, detail string) {
	HTTPRespondError(w, 401, "unauthorized", detail)
}

func HTTPRespond403(w http.ResponseWriter, detail string) {
	HTTPRespondError(w, 403, "permission_denied", detail)
}

func HTTPRespond404(w http.ResponseWriter, detail string) {
	HTTPRespondError(w, 404, "not_found", detail)
}
<commit_msg>upgrade http<commit_after>package lily

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	netUrl "net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

func HTTPStatusCodeIsOk(statusCode int) bool {
	return statusCode > 199 && statusCode < 300
}

func HTTPSetContentTypeJSON(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
}

func HTTPSetContentTypeHTML(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/html; charset=UTF-8")
}

func HTTPRespondStr(w http.ResponseWriter, code int, body string) {
	if len(body) == 0 {
		panic("body must be not empty")
	}
	w.WriteHeader(code)
	fmt.Fprint(w, body)
}

func HTTPRespondJSONObj(w http.ResponseWriter, code int, obj interface{}) {
	HTTPSetContentTypeJSON(w)
	w.WriteHeader(code)
	ErrPanic(json.NewEncoder(w).Encode(obj))
}

func HTTPRespondJSONParseError(w http.ResponseWriter) {
	HTTPRespond400(w, "bad_json", "Fail to parse JSON")
}

func HTTPSendRequest(withJar bool, method, url string, urlParams map[string]string,
	data []byte, timeout time.Duration, headers ...string) (*http.Response, error) {
	var err error
	var req *http.Request
	var jar http.CookieJar

	if data != nil {
		req, err = http.NewRequest(method, url, bytes.NewBuffer(data))
		ErrPanic(err)
	} else {
		req, err = http.NewRequest(method, url, nil)
		ErrPanic(err)
	}

	if urlParams != nil {
		q := netUrl.Values{}
		for k, v := range urlParams {
			q.Add(k, v)
		}
		req.URL.RawQuery = q.Encode()
	}

	for i := 0; (i + 1) < len(headers); i += 2 {
		req.Header.Set(headers[i], headers[i+1])
	}

	if withJar {
		jar, err = cookiejar.New(nil)
		ErrPanic(err)
	}

	client := http.Client{
		Timeout: timeout,
		Jar:     jar,
	}

	return client.Do(req)
}

func HTTPSendRequestReceiveBytes(withJar, errSCode bool, method, url string, urlParams map[string]string,
	data []byte, timeout time.Duration, headers ...string) (int, []byte, error) {
	var res []byte

	resp, err := HTTPSendRequest(withJar, method, url, urlParams, data, timeout, headers...)
	if err != nil {
		return 0, nil, err
	}
	defer resp.Body.Close()

	res, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0, nil, err
	}

	if !HTTPStatusCodeIsOk(resp.StatusCode) {
		if errSCode {
			return resp.StatusCode, res, errors.New(fmt.Sprintf("bad_http_status_code - %d\nbody: %s", resp.StatusCode, string(res)))
		}
		return resp.StatusCode, res, nil
	}

	return resp.StatusCode, res, nil
}

func HTTPSendRequestReceiveString(withJar, errSCode bool, method, url string, urlParams map[string]string,
	data []byte, timeout time.Duration, headers ...string) (int, string, error) {
	sCode, resBytes, err := HTTPSendRequestReceiveBytes(withJar, errSCode, method, url, urlParams, data, timeout, headers...)

	return sCode, string(resBytes), err
}

func HTTPSendRequestReceiveJSONObj(withJar, errSCode bool, method, url string, urlParams map[string]string,
	data []byte, rObj interface{}, timeout time.Duration, headers ...string) (int, []byte, error) {
	sCode, rBytes, err := HTTPSendRequestReceiveBytes(
		withJar, errSCode, method, url, urlParams, data, timeout, headers...)
	if err != nil || !HTTPStatusCodeIsOk(sCode) {
		return sCode, rBytes, err
	}

	err = json.Unmarshal(rBytes, rObj)
	if err != nil {
		return sCode, rBytes, errors.New(fmt.Sprintf("fail_to_parse_json - %s\nbody: %s", err.Error(), string(rBytes)))
	}

	return sCode, rBytes, nil
}

func HTTPSendJSONRequestReceiveBytes(withJar, errSCode bool, method, url string, urlParams map[string]string,
	sObj interface{}, timeout time.Duration, headers ...string) (int, []byte, error) {
	sBytes, err := json.Marshal(sObj)
	if err != nil {
		return 0, nil, err
	}

	return HTTPSendRequestReceiveBytes(withJar, errSCode, method, url, urlParams, sBytes, timeout, headers...)
}

func HTTPSendJSONRequestReceiveString(withJar, errSCode bool, method, url string, urlParams map[string]string,
	sObj interface{}, timeout time.Duration, headers ...string) (int, string, error) {
	sBytes, err := json.Marshal(sObj)
	if err != nil {
		return 0, "", err
	}

	return HTTPSendRequestReceiveString(withJar, errSCode, method, url, urlParams, sBytes, timeout, headers...)
}

func HTTPSendJSONRequestReceiveJSONObj(withJar, errSCode bool, method, url string, urlParams map[string]string,
	sObj interface{}, rObj interface{}, timeout time.Duration, headers ...string) (int, []byte, error) {
	sBytes, err := json.Marshal(sObj)
	if err != nil {
		return 0, nil, err
	}

	return HTTPSendRequestReceiveJSONObj(withJar, errSCode, method, url, urlParams, sBytes, rObj, timeout, headers...)
}

func HTTPRetrieveRequestHostURL(r *http.Request) string {
	scheme := r.Header.Get("X-Forwarded-Proto")
	if scheme == "" {
		if r.TLS == nil {
			scheme = "http"
		} else {
			scheme = "https"
		}
	}
	return scheme + "://" + r.Host
}

func HTTPRetrieveRemoteIP(r *http.Request) (result string) {
	result = ""
	if parts := strings.Split(r.RemoteAddr, ":"); len(parts) == 2 {
		result = parts[0]
	}
	// If we have a forwarded-for header, take the address from there
	if xff := strings.Trim(r.Header.Get("X-Forwarded-For"), ","); len(xff) > 0 {
		addrs := strings.Split(xff, ",")
		lastFwd := addrs[len(addrs)-1]
		if ip := net.ParseIP(lastFwd); ip != nil {
			result = ip.String()
		}
		// parse X-Real-Ip header
	} else if xri := r.Header.Get("X-Real-Ip"); len(xri) > 0 {
		if ip := net.ParseIP(xri); ip != nil {
			result = ip.String()
		}
	}
	return
}

func HTTPUploadFileFromRequestForm(r *http.Request, key, dirPath, dir string, filename string) (string, error) {
	var err error

	finalDirPath := filepath.Join(dirPath, dir)

	err = os.MkdirAll(finalDirPath, os.ModePerm)
	if err != nil {
		return "", err
	}

	srcFile, header, err := r.FormFile(key)
	if err != nil {
		return "", err
	}
	defer srcFile.Close()

	fileExt := filepath.Ext(header.Filename)
	if fileExt == "" {
		return "", errors.New("bad_extension")
	}

	dstFile, err := TempFile(finalDirPath, filename+"_*"+fileExt)
	if err != nil {
		return "", err
	}
	defer dstFile.Close()

	_, err = io.Copy(dstFile, srcFile)
	if err != nil {
		return "", err
	}

	err = os.Chmod(dstFile.Name(), 0644)
	if err != nil {
		return "", err
	}

	newName, err := filepath.Rel(dirPath, dstFile.Name())
	if err != nil {
		return "", err
	}

	return newName, nil
}

func HTTPRespondError(w http.ResponseWriter, code int, err string, detail string, extras ...interface{}) {
	obj := map[string]interface{}{}
	obj["error"] = err
	obj["error_dsc"] = detail
	for i := 0; (i + 1) < len(extras); i += 2 {
		obj[extras[i].(string)] = extras[i+1]
	}
	HTTPRespondJSONObj(w, code, obj)
}

func HTTPRespond400(w http.ResponseWriter, err, detail string, extras ...interface{}) {
	HTTPRespondError(w, 400, err, detail, extras...)
}

func HTTPRespond401(w http.ResponseWriter, detail string) {
	HTTPRespondError(w, 401, "unauthorized", detail)
}

func HTTPRespond403(w http.ResponseWriter, detail string) {
	HTTPRespondError(w, 403, "permission_denied", detail)
}

func HTTPRespond404(w http.ResponseWriter, detail string) {
	HTTPRespondError(w, 404, "not_found", detail)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"time"
)

type HTTPServer struct {
	listenAddr string
	vdt        *VDirTree
}

func (svr *HTTPServer) startServer() {
	svr.vdt = &VDirTree{
		rpath: "",
		vpath: "/",
	}
	http.HandleFunc("/", svr.rootHandle)
	http.HandleFunc("/static/", svr.staticHandle)
	http.HandleFunc("/fs/", svr.fsHandle)
	http.ListenAndServe(svr.listenAddr, nil)
}

func (svr *HTTPServer) shutdownServer() {
}

func (svr *HTTPServer) rootHandle(w http.ResponseWriter, r *http.Request) {
	http.Redirect(w, r, "/fs/", 308)
}

func (svr *HTTPServer) staticHandle(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, r.URL.Path[1:])
}

func humanizeSize(s int64) string {
	size := float64(s)
	units := []string{"B", "KB", "MB", "GB", "TB"}
	for i, u := range units {
		if size < 9*1000.0 {
			if i == 0 {
				break
			}
			return fmt.Sprintf("%.02f %s", size, u)
		}
		size /= 1024.0
	}
	return fmt.Sprintf("%d B", int64(s))
}

func humanizeTime(t time.Time) string {
	return fmt.Sprintf("%02d-%s-%04d %02d:%02d",
		t.Day(), t.Month().String()[:3], t.Year(),
		t.Hour(), t.Minute(),
	)
}

func pathHiSplit(p string) []string {
	dir := path.Dir(p)
	base := path.Base(p)
	if dir == base {
		return []string{"/"}
	}
	return append(pathHiSplit(dir), p)
}

var funcMap = template.FuncMap{
	"humanize_size": humanizeSize,
	"humanize_time": humanizeTime,
	"path_base":     func(p string) string { return path.Base(p) },
	"path_dir":      func(p string) string { return path.Dir(p) },
	"path_hisplit":  pathHiSplit,
	"plus1":         func(x int) int { return x + 1 },
}

func (svr *HTTPServer) fsHandle(w http.ResponseWriter, r *http.Request) {
	vpath := path.Join("/", r.URL.Path[len("/fs/"):])
	rp, vp, err := svr.vdt.mapVPath(vpath)
	if err != nil {
		http.Error(w, http.StatusText(404), 404)
		return
	}
	switch r.Method {
	case "GET":
		svr.getHandle(rp, vp, w, r)
	case "POST":
		if !conf.writable {
			http.Error(w, http.StatusText(403), 403)
			return
		}
		svr.putHandle(rp, vp, w, r)
	}
}

func (svr *HTTPServer) putHandle(rp, vp string, w http.ResponseWriter, r *http.Request) {
	err := r.ParseMultipartForm(200000)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	formdata := r.MultipartForm
	for _, fhead := range formdata.File {
		fh, fn := fhead[0], fhead[0].Filename
		file, err := fh.Open()
		defer file.Close()
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}

		out, err := os.Create(filepath.Join(rp, fn))
		defer out.Close()
		if err != nil {
			http.Error(w, "Unable to create the file for writing.", 500)
			return
		}

		log.Printf("Receiving file: %s", fn)
		_, err = io.Copy(out, file)
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
	}
}

func (svr *HTTPServer) getHandle(rp, vp string, w http.ResponseWriter, r *http.Request) {
	if vp != "/" {
		fi, err := os.Stat(rp)
		if err != nil {
			http.Error(w, http.StatusText(500), 500)
			return
		}
		if fi.Mode().IsRegular() {
			http.ServeFile(w, r, rp)
			return
		}
	}
	lst, err := svr.vdt.doLIST(vp)
	if err != nil {
		http.Error(w, http.StatusText(500), 500)
		return
	}
	tpl := template.Must(template.New("main").Funcs(funcMap).ParseGlob("*.html"))
	content := map[string]interface{}{
		"CWD":        path.Clean(vp),
		"IsVRoot":    vp == "/",
		"Items":      lst,
		"IsWritable": conf.writable && vp != "/",
	}
	tpl.ExecuteTemplate(w, "index.html", content)
}
<commit_msg>http server can show very large file size.<commit_after>package main

import (
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"time"
)

type HTTPServer struct {
	listenAddr string
	vdt        *VDirTree
}

func (svr *HTTPServer) startServer() {
	svr.vdt = &VDirTree{
		rpath: "",
		vpath: "/",
	}
	http.HandleFunc("/", svr.rootHandle)
	http.HandleFunc("/static/", svr.staticHandle)
	http.HandleFunc("/fs/", svr.fsHandle)
	http.ListenAndServe(svr.listenAddr, nil)
}

func (svr *HTTPServer) shutdownServer() {
}

func (svr *HTTPServer) rootHandle(w http.ResponseWriter, r *http.Request) {
	http.Redirect(w, r, "/fs/", 308)
}

func (svr *HTTPServer) staticHandle(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, r.URL.Path[1:])
}

func humanizeSize(s int64) string {
	size := float64(s)
	units := []string{"B", "KB", "MB", "GB", "TB", "PB", "EB"}
	for i, u := range units {
		if size < 9*1000.0 {
			if i == 0 {
				break
			}
			return fmt.Sprintf("%.02f %s", size, u)
		}
		size /= 1024.0
	}
	return fmt.Sprintf("%d B", int64(s))
}

func humanizeTime(t time.Time) string {
	return fmt.Sprintf("%02d-%s-%04d %02d:%02d",
		t.Day(), t.Month().String()[:3], t.Year(),
		t.Hour(), t.Minute(),
	)
}

func pathHiSplit(p string) []string {
	dir := path.Dir(p)
	base := path.Base(p)
	if dir == base {
		return []string{"/"}
	}
	return append(pathHiSplit(dir), p)
}

var funcMap = template.FuncMap{
	"humanize_size": humanizeSize,
	"humanize_time": humanizeTime,
	"path_base":     func(p string) string { return path.Base(p) },
	"path_dir":      func(p string) string { return path.Dir(p) },
	"path_hisplit":  pathHiSplit,
	"plus1":         func(x int) int { return x + 1 },
}

func (svr *HTTPServer) fsHandle(w http.ResponseWriter, r *http.Request) {
	vpath := path.Join("/", r.URL.Path[len("/fs/"):])
	rp, vp, err := svr.vdt.mapVPath(vpath)
	if err != nil {
		http.Error(w, http.StatusText(404), 404)
		return
	}
	switch r.Method {
	case "GET":
		svr.getHandle(rp, vp, w, r)
	case "POST":
		if !conf.writable {
			http.Error(w, http.StatusText(403), 403)
			return
		}
		svr.putHandle(rp, vp, w, r)
	}
}

func (svr *HTTPServer) putHandle(rp, vp string, w http.ResponseWriter, r *http.Request) {
	err := r.ParseMultipartForm(200000)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	formdata := r.MultipartForm
	for _, fhead := range formdata.File {
		fh, fn := fhead[0], fhead[0].Filename
		file, err := fh.Open()
		defer file.Close()
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}

		out, err := os.Create(filepath.Join(rp, fn))
		defer out.Close()
		if err != nil {
			http.Error(w, "Unable to create the file for writing.", 500)
			return
		}

		log.Printf("Receiving file: %s", fn)
		_, err = io.Copy(out, file)
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
	}
}

func (svr *HTTPServer) getHandle(rp, vp string, w http.ResponseWriter, r *http.Request) {
	if vp != "/" {
		fi, err := os.Stat(rp)
		if err != nil {
			http.Error(w, http.StatusText(500), 500)
			return
		}
		if fi.Mode().IsRegular() {
			http.ServeFile(w, r, rp)
			return
		}
	}
	lst, err := svr.vdt.doLIST(vp)
	if err != nil {
		http.Error(w, http.StatusText(500), 500)
		return
	}
	tpl := template.Must(template.New("main").Funcs(funcMap).ParseGlob("*.html"))
	content := map[string]interface{}{
		"CWD":        path.Clean(vp),
		"IsVRoot":    vp == "/",
		"Items":      lst,
		"IsWritable": conf.writable && vp != "/",
	}
	tpl.ExecuteTemplate(w, "index.html", content)
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Eric Barkie. All rights reserved.
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package main

// HTTP server for accessing weather station data.

import (
	"encoding/json"
	"fmt"
	"net/http"
	//_ "net/http/pprof"
	"strconv"
	"time"
)

type httpContext serverContext

type httpLogWrapper struct {
	http.CloseNotifier
	http.Flusher
	http.ResponseWriter
	status int
}

func (l *httpLogWrapper) Write(p []byte) (int, error) {
	return l.ResponseWriter.Write(p)
}

func (l *httpLogWrapper) WriteHeader(status int) {
	l.status = status
	l.ResponseWriter.WriteHeader(status)
}

func (httpContext) logHandler(h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		record := &httpLogWrapper{
			CloseNotifier:  w.(http.CloseNotifier),
			Flusher:        w.(http.Flusher),
			ResponseWriter: w,
			status:         http.StatusOK,
		}

		h.ServeHTTP(record, r)

		msg := fmt.Sprintf("HTTP connection from %s request %s %s response %d", r.RemoteAddr, r.Method, r.URL, record.status)
		if record.status < 299 {
			Debug.Print(msg)
		} else {
			Warn.Print(msg)
		}
	}
}

// archive is the endpoint for serving out archive records.
// GET /archive[?begin=2016-08-03T00:00:00Z][&end=2016-09-03T00:00:00Z]
func (c httpContext) archive(w http.ResponseWriter, r *http.Request) {
	// Parse and validate begin and end parameters.
	var begin, end time.Time
	var err error

	if r.URL.Query().Get("end") != "" {
		end, err = time.Parse(time.RFC3339, r.URL.Query().Get("end"))
		if err != nil {
			w.Header().Set("Warning", "Unable to parse end timestamp")
			w.WriteHeader(http.StatusBadRequest)
			return
		}
	} else {
		// Default end is now.
		end = time.Now()
	}

	if r.URL.Query().Get("begin") != "" {
		begin, err = time.Parse(time.RFC3339, r.URL.Query().Get("begin"))
		if err != nil {
			w.Header().Set("Warning", "Unable to parse begin timestamp")
			w.WriteHeader(http.StatusBadRequest)
			return
		}
	} else {
		// Default begin is 1 day before end.
		begin = end.AddDate(0, 0, -1)
	}

	if end.Before(begin) {
		w.Header().Set("Warning", "End timestamp precedes begin timestamp")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Large durations can be very resource intensive to marshal so
	// cap at 30 days.
	if end.Sub(begin) > (30 * (24 * time.Hour)) {
		w.Header().Set("Warning", "Duration exceeds maximum allowed")
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		return
	}

	// Query archive from database and return.
	archive := c.ad.Get(begin, end)
	if len(archive) < 1 {
		w.WriteHeader(http.StatusNoContent)
		return
	}
	j, _ := json.MarshalIndent(archive, "", "    ")
	w.Header().Set("Content-Type", "application/json")
	w.Write(j)
}

// loop is the endpoint for serving out loop samples.
// GET /loop[?lastSequence=#]
func (c httpContext) loop(w http.ResponseWriter, r *http.Request) {
	c.ld.RLock()
	defer c.ld.RUnlock()

	// If there aren't enough samples (the server just started) or
	// there were no recent updates then send a HTTP service temporarily
	// unavailable response.
	if len(c.ld.loops) < loopsMin {
		w.Header().Set("Warning", "Not enough samples yet")
		w.WriteHeader(http.StatusServiceUnavailable)
	} else if time.Since(c.ld.loops[0].Update.Timestamp) > loopStaleAge {
		w.Header().Set("Warning", "Samples are too old")
		w.WriteHeader(http.StatusServiceUnavailable)
	} else {
		// Figure out if request is for loops since a sequence or just for
		// most recent loop.
		var j []byte
		if r.URL.Query().Get("lastSequence") != "" {
			seq, _ := strconv.ParseInt(r.URL.Query().Get("lastSequence"), 10, 64)

			// There are no sequence gaps so it's simple subtraction to
			// determine the end index.  A few safeguards have to be added
			// though:
			//
			// If the requested sequence is ahead of the server then return
			// nothing.
			//
			// If the request sequence is so far back that it's been purged
			// then return everything.
			endIndex := int(c.ld.loops[0].Update.Sequence - seq)
			if endIndex < 1 {
				j, _ = json.Marshal(nil)
			} else {
				if endIndex > len(c.ld.loops) {
					endIndex = len(c.ld.loops)
				}
				j, _ = json.MarshalIndent(c.ld.loops[0:endIndex], "", "    ")
			}
		} else {
			j, _ = json.MarshalIndent(c.ld.loops[0], "", "    ")
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(j)
	}
}

// events is the endpoint for streaming loop samples using the Server-sent
// events.
// GET /events
func (c httpContext) events(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	incomingEvents := c.eb.subscribe(r.RemoteAddr)
	defer c.eb.unsubscribe(incomingEvents)

	for {
		select {
		case <-w.(http.CloseNotifier).CloseNotify():
			// Client closed the connection
			return
		case e := <-incomingEvents:
			fmt.Fprintf(w, "event: %s\n", e.event)
			r, _ := json.Marshal(e.data)
			fmt.Fprintf(w, "data: %s\n\n", r)
			w.(http.Flusher).Flush()
		}
	}
}

// httpServer starts the HTTP server.  It's blocking and should be called as
// a goroutine.
func httpServer(bindAddress string, sc serverContext) {
	c := httpContext(sc)
	http.HandleFunc("/archive", c.archive)
	http.HandleFunc("/loop", c.loop)
	http.HandleFunc("/events", c.events)

	s := http.Server{
		Addr:    bindAddress + ":8080",
		Handler: c.logHandler(http.DefaultServeMux),
	}
	Info.Printf("HTTP server started on %s", s.Addr)
	err := s.ListenAndServe()
	if err != nil {
		Error.Fatalf("HTTP server error: %s", err.Error())
	}
}
<commit_msg>Remove excessive comments<commit_after>// Copyright (c) 2016-2017 Eric Barkie. All rights reserved.
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package main

// HTTP server for accessing weather station data.

import (
	"encoding/json"
	"fmt"
	"net/http"
	//_ "net/http/pprof"
	"strconv"
	"time"
)

type httpContext serverContext

type httpLogWrapper struct {
	http.CloseNotifier
	http.Flusher
	http.ResponseWriter
	status int
}

func (l *httpLogWrapper) Write(p []byte) (int, error) {
	return l.ResponseWriter.Write(p)
}

func (l *httpLogWrapper) WriteHeader(status int) {
	l.status = status
	l.ResponseWriter.WriteHeader(status)
}

func (httpContext) logHandler(h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		record := &httpLogWrapper{
			CloseNotifier:  w.(http.CloseNotifier),
			Flusher:        w.(http.Flusher),
			ResponseWriter: w,
			status:         http.StatusOK,
		}

		h.ServeHTTP(record, r)

		msg := fmt.Sprintf("HTTP connection from %s request %s %s response %d", r.RemoteAddr, r.Method, r.URL, record.status)
		if record.status < 299 {
			Debug.Print(msg)
		} else {
			Warn.Print(msg)
		}
	}
}

// archive is the endpoint for serving out archive records.
// GET /archive[?begin=2016-08-03T00:00:00Z][&end=2016-09-03T00:00:00Z]
func (c httpContext) archive(w http.ResponseWriter, r *http.Request) {
	// Parse and validate begin and end parameters.
	var begin, end time.Time
	var err error

	if r.URL.Query().Get("end") != "" {
		end, err = time.Parse(time.RFC3339, r.URL.Query().Get("end"))
		if err != nil {
			w.Header().Set("Warning", "Unable to parse end timestamp")
			w.WriteHeader(http.StatusBadRequest)
			return
		}
	} else {
		end = time.Now()
	}

	if r.URL.Query().Get("begin") != "" {
		begin, err = time.Parse(time.RFC3339, r.URL.Query().Get("begin"))
		if err != nil {
			w.Header().Set("Warning", "Unable to parse begin timestamp")
			w.WriteHeader(http.StatusBadRequest)
			return
		}
	} else {
		begin = end.AddDate(0, 0, -1)
	}

	if end.Before(begin) {
		w.Header().Set("Warning", "End timestamp precedes begin timestamp")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Large durations can be very resource intensive to marshal so
	// cap at 30 days.
	if end.Sub(begin) > (30 * (24 * time.Hour)) {
		w.Header().Set("Warning", "Duration exceeds maximum allowed")
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		return
	}

	// Query archive from database and return.
	archive := c.ad.Get(begin, end)
	if len(archive) < 1 {
		w.WriteHeader(http.StatusNoContent)
		return
	}
	j, _ := json.MarshalIndent(archive, "", "    ")
	w.Header().Set("Content-Type", "application/json")
	w.Write(j)
}

// loop is the endpoint for serving out loop samples.
// GET /loop[?lastSequence=#]
func (c httpContext) loop(w http.ResponseWriter, r *http.Request) {
	c.ld.RLock()
	defer c.ld.RUnlock()

	// If there aren't enough samples (the server just started) or
	// there were no recent updates then send a HTTP service temporarily
	// unavailable response.
	if len(c.ld.loops) < loopsMin {
		w.Header().Set("Warning", "Not enough samples yet")
		w.WriteHeader(http.StatusServiceUnavailable)
	} else if time.Since(c.ld.loops[0].Update.Timestamp) > loopStaleAge {
		w.Header().Set("Warning", "Samples are too old")
		w.WriteHeader(http.StatusServiceUnavailable)
	} else {
		// Figure out if request is for loops since a sequence or just for
		// most recent loop.
		var j []byte
		if r.URL.Query().Get("lastSequence") != "" {
			seq, _ := strconv.ParseInt(r.URL.Query().Get("lastSequence"), 10, 64)

			// There are no sequence gaps so it's simple subtraction to
			// determine the end index.  A few safeguards have to be added
			// though:
			//
			// If the requested sequence is ahead of the server then return
			// nothing.
			//
			// If the request sequence is so far back that it's been purged
			// then return everything.
			endIndex := int(c.ld.loops[0].Update.Sequence - seq)
			if endIndex < 1 {
				j, _ = json.Marshal(nil)
			} else {
				if endIndex > len(c.ld.loops) {
					endIndex = len(c.ld.loops)
				}
				j, _ = json.MarshalIndent(c.ld.loops[0:endIndex], "", "    ")
			}
		} else {
			j, _ = json.MarshalIndent(c.ld.loops[0], "", "    ")
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(j)
	}
}

// events is the endpoint for streaming loop samples using the Server-sent
// events.
// GET /events
func (c httpContext) events(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	incomingEvents := c.eb.subscribe(r.RemoteAddr)
	defer c.eb.unsubscribe(incomingEvents)

	for {
		select {
		case <-w.(http.CloseNotifier).CloseNotify():
			// Client closed the connection
			return
		case e := <-incomingEvents:
			fmt.Fprintf(w, "event: %s\n", e.event)
			r, _ := json.Marshal(e.data)
			fmt.Fprintf(w, "data: %s\n\n", r)
			w.(http.Flusher).Flush()
		}
	}
}

// httpServer starts the HTTP server.  It's blocking and should be called as
// a goroutine.
func httpServer(bindAddress string, sc serverContext) {
	c := httpContext(sc)
	http.HandleFunc("/archive", c.archive)
	http.HandleFunc("/loop", c.loop)
	http.HandleFunc("/events", c.events)

	s := http.Server{
		Addr:    bindAddress + ":8080",
		Handler: c.logHandler(http.DefaultServeMux),
	}
	Info.Printf("HTTP server started on %s", s.Addr)
	err := s.ListenAndServe()
	if err != nil {
		Error.Fatalf("HTTP server error: %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/flynn/go-discoverd"
	"github.com/flynn/strowger/types"
	"github.com/inconshreveable/go-vhost"
)

type HTTPListener struct {
	Watcher

	Addr      string
	TLSAddr   string
	TLSConfig *tls.Config

	mtx      sync.RWMutex
	domains  map[string]*httpRoute
	services map[string]*httpService

	discoverd DiscoverdClient
	ds        DataStore
	wm        *WatchManager

	listener    net.Listener
	tlsListener net.Listener
	closed      bool
}

type DiscoverdClient interface {
	NewServiceSet(string) (discoverd.ServiceSet, error)
}

func NewHTTPListener(addr, tlsAddr string, ds DataStore, discoverdc DiscoverdClient) *HTTPListener {
	l := &HTTPListener{
		Addr:      addr,
		TLSAddr:   tlsAddr,
		ds:        ds,
		discoverd: discoverdc,
		domains:   make(map[string]*httpRoute),
		services:  make(map[string]*httpService),
		wm:        NewWatchManager(),
	}
	l.Watcher = l.wm
	return l
}

func (s *HTTPListener) Close() error {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	for _, service := range s.services {
		service.ss.Close()
	}
	s.listener.Close()
	s.tlsListener.Close()
	s.ds.StopSync()
	s.closed = true
	return nil
}

func (s *HTTPListener) Start() error {
	started := make(chan error)

	go s.ds.Sync(&httpSyncHandler{l: s}, started)
	if err := <-started; err != nil {
		return err
	}

	go s.serve(started)
	if err := <-started; err != nil {
		s.ds.StopSync()
		return err
	}
	s.Addr = s.listener.Addr().String()

	go s.serveTLS(started)
	if err := <-started; err != nil {
		s.ds.StopSync()
		s.listener.Close()
		return err
	}
	s.TLSAddr = s.tlsListener.Addr().String()

	return nil
}

var ErrClosed = errors.New("strowger: listener has been closed")

func (s *HTTPListener) AddRoute(r *strowger.HTTPRoute) error {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	if s.closed {
		return ErrClosed
	}

	return s.ds.Add(r.Domain, &httpRoute{
		Domain:  r.Domain,
		Service: r.Service,
		TLSCert: r.TLSCert,
		TLSKey:  r.TLSKey,
	})
}

func (s *HTTPListener) RemoveRoute(domain string) error {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	if s.closed {
		return ErrClosed
	}

	return s.ds.Remove(domain)
}

type httpSyncHandler struct {
	l *HTTPListener
}

func (h *httpSyncHandler) Add(data []byte) error {
	r := &httpRoute{}
	if err := json.Unmarshal(data, r); err != nil {
		return err
	}

	if r.TLSCert != "" && r.TLSKey != "" {
		kp, err := tls.X509KeyPair([]byte(r.TLSCert), []byte(r.TLSKey))
		if err != nil {
			return err
		}
		r.keypair = &kp
		r.TLSCert = ""
		r.TLSKey = ""
	}

	h.l.mtx.Lock()
	defer h.l.mtx.Unlock()
	if h.l.closed {
		return ErrClosed
	}
	if _, ok := h.l.domains[r.Domain]; ok {
		return ErrExists
	}

	service := h.l.services[r.Service]
	if service == nil {
		ss, err := h.l.discoverd.NewServiceSet(r.Service)
		if err != nil {
			return err
		}
		service = &httpService{name: r.Service, ss: ss}
		h.l.services[r.Service] = service
	}
	service.refs++
	r.service = service
	h.l.domains[r.Domain] = r

	log.Println("Adding domain", r.Domain)
	go h.l.wm.Send(&strowger.Event{Event: "add", ID: r.Domain})
	return nil
}

func (h *httpSyncHandler) Remove(name string) error {
	h.l.mtx.Lock()
	defer h.l.mtx.Unlock()
	r, ok := h.l.domains[name]
	if !ok {
		return ErrNotFound
	}

	r.service.refs--
	if r.service.refs <= 0 {
		r.service.ss.Close()
		delete(h.l.services, r.service.name)
	}

	delete(h.l.domains, name)
	log.Println("Removing domain", name)
	go h.l.wm.Send(&strowger.Event{Event: "remove", ID: name})
	return nil
}

func (s *HTTPListener) serve(started chan<- error) {
	var err error
	s.listener, err = net.Listen("tcp", s.Addr)
	started <- err
	if err != nil {
		return
	}
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			// TODO: log error
			break
		}
		go s.handle(conn, false)
	}
}

func (s *HTTPListener) serveTLS(started chan<- error) {
	var err error
	s.tlsListener, err = net.Listen("tcp", s.TLSAddr)
	started <- err
	if err != nil {
		return
	}
	for {
		conn, err := s.tlsListener.Accept()
		if err != nil {
			// TODO: log error
			break
		}
		go s.handle(conn, true)
	}
}

func (s *HTTPListener) findRouteForHost(host string) *httpRoute {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	// TODO: handle wildcard domains
	backend := s.domains[host]
	log.Printf("Backend match: %#v\n", backend)
	return backend
}

func fail(sc *httputil.ServerConn, req *http.Request, code int, msg string) {
	resp := &http.Response{
		StatusCode:    code,
		ProtoMajor:    1,
		ProtoMinor:    0,
		Request:       req,
		Body:          ioutil.NopCloser(bytes.NewBufferString(msg)),
		ContentLength: int64(len(msg)),
	}
	sc.Write(req, resp)
}

func (s *HTTPListener) handle(conn net.Conn, isTLS bool) {
	defer conn.Close()

	var r *httpRoute

	// For TLS, use the SNI hello to determine the domain.
	// At this stage, if we don't find a match, we simply
	// close the connection down.
	if isTLS {
		// Parse out host via SNI first
		vhostConn, err := vhost.TLS(conn)
		if err != nil {
			log.Println("Failed to decode TLS connection", err)
			return
		}
		host := vhostConn.Host()
		log.Println("SNI host is:", host)

		// Find a backend for the key
		r = s.findRouteForHost(host)
		if r == nil {
			return
		}
		if r.keypair == nil {
			log.Println("Cannot serve TLS, no certificate defined for this domain")
			return
		}

		// Init a TLS decryptor
		tlscfg := &tls.Config{Certificates: []tls.Certificate{*r.keypair}}
		conn = tls.Server(vhostConn, tlscfg)
	}

	// Decode the first request from the connection
	sc := httputil.NewServerConn(conn, nil)
	req, err := sc.Read()
	if err != nil {
		if err != httputil.ErrPersistEOF {
			// TODO: log error
		}
		return
	}

	// If we do not have a backend yet (unencrypted connection),
	// look at the host header to find one or 404 out.
	if r == nil {
		r = s.findRouteForHost(req.Host)
		if r == nil {
			fail(sc, req, 404, "Not Found")
			return
		}
	}

	r.service.handle(req, sc, isTLS)
}

// A domain served by a listener, associated TLS certs,
// and link to backend service set.
type httpRoute struct {
	Domain  string
	Service string
	TLSCert string
	TLSKey  string

	keypair *tls.Certificate
	service *httpService
}

// A service definition: name, and set of backends.
type httpService struct {
	name string
	ss   discoverd.ServiceSet
	refs int
}

func (s *httpService) getBackend() *httputil.ClientConn {
	for _, addr := range shuffle(s.ss.Addrs()) {
		// TODO: set connection timeout
		backend, err := net.Dial("tcp", addr)
		if err != nil {
			// TODO: log error
			// TODO: limit number of backends tried
			// TODO: temporarily quarantine failing backends
			log.Println("backend error", err)
			continue
		}
		return httputil.NewClientConn(backend, nil)
	}
	// TODO: log no backends found error
	return nil
}

func (s *httpService) handle(req *http.Request, sc *httputil.ServerConn, tls bool) {
	req.Header.Set("X-Request-Start", strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10))
	backend := s.getBackend()
	if backend == nil {
		log.Println("no backend found")
		fail(sc, req, 503, "Service Unavailable")
		return
	}
	defer backend.Close()

	for {
		if req.Method != "GET" && req.Method != "POST" && req.Method != "HEAD" &&
			req.Method != "OPTIONS" && req.Method != "PUT" && req.Method != "DELETE" && req.Method != "TRACE" {
			fail(sc, req, 405, "Method not allowed")
			return
		}

		req.Proto = "HTTP/1.1"
		req.ProtoMajor = 1
		req.ProtoMinor = 1
		delete(req.Header, "Te")
		delete(req.Header, "Transfer-Encoding")

		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			// If we aren't the first proxy retain prior
			// X-Forwarded-For information as a comma+space
			// separated list and fold multiple headers into one.
			if prior, ok := req.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			req.Header.Set("X-Forwarded-For", clientIP)
		}
		if tls {
			req.Header.Set("X-Forwarded-Proto", "https")
		} else {
			req.Header.Set("X-Forwarded-Proto", "http")
		}
		// TODO: Set X-Forwarded-Port

		if err := backend.Write(req); err != nil {
			log.Println("server write err:", err)
			return
		}
		res, err := backend.Read(req)
		if res != nil {
			if err := sc.Write(req, res); err != nil {
				if err != io.EOF && err != httputil.ErrPersistEOF {
					log.Println("client write err:", err)
					// TODO: log error
				}
				return
			}
		}
		if err != nil {
			if err != io.EOF && err != httputil.ErrPersistEOF {
				log.Println("server read err:", err)
				// TODO: log error
				fail(sc, req, 502, "Bad Gateway")
			}
			return
		}

		// TODO: Proxy HTTP CONNECT? (example: Go RPC over HTTP)
		if res.StatusCode == http.StatusSwitchingProtocols {
			serverW, serverR := backend.Hijack()
			clientW, clientR := sc.Hijack()
			defer serverW.Close()
			done := make(chan struct{})
			go func() {
				serverR.WriteTo(clientW)
				close(done)
			}()
			clientR.WriteTo(serverW)
			<-done
			return
		}

		// TODO: http pipelining
		req, err = sc.Read()
		if err != nil {
			if err != io.EOF && err != httputil.ErrPersistEOF {
				log.Println("client read err:", err)
			}
			return
		}
		req.Header.Set("X-Request-Start", strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10))
	}
}

func shuffle(s []string) []string {
	for i := len(s) - 1; i > 0; i-- {
		j := rand.Intn(i + 1)
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func init() {
	rand.Seed(time.Now().UnixNano())
}
<commit_msg>router: Remove verbose logging in non-error paths<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/flynn/go-discoverd"
	"github.com/flynn/strowger/types"
	"github.com/inconshreveable/go-vhost"
)

type HTTPListener struct {
	Watcher

	Addr      string
	TLSAddr   string
	TLSConfig *tls.Config

	mtx      sync.RWMutex
	domains  map[string]*httpRoute
	services map[string]*httpService

	discoverd DiscoverdClient
	ds        DataStore
	wm        *WatchManager

	listener    net.Listener
	tlsListener net.Listener
	closed      bool
}

type DiscoverdClient interface {
	NewServiceSet(string) (discoverd.ServiceSet, error)
}

func NewHTTPListener(addr, tlsAddr string, ds DataStore, discoverdc DiscoverdClient) *HTTPListener {
	l := &HTTPListener{
		Addr:      addr,
		TLSAddr:   tlsAddr,
		ds:        ds,
		discoverd: discoverdc,
		domains:   make(map[string]*httpRoute),
		services:  make(map[string]*httpService),
		wm:        NewWatchManager(),
	}
	l.Watcher = l.wm
	return l
}

func (s *HTTPListener) Close() error {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	for _, service := range s.services {
		service.ss.Close()
	}
	s.listener.Close()
	s.tlsListener.Close()
	s.ds.StopSync()
	s.closed = true
	return nil
}

func (s *HTTPListener) Start() error {
	started := make(chan error)

	go s.ds.Sync(&httpSyncHandler{l: s}, started)
	if err := <-started; err != nil {
		return err
	}

	go s.serve(started)
	if err := <-started; err != nil {
		s.ds.StopSync()
		return err
	}
	s.Addr = s.listener.Addr().String()

	go s.serveTLS(started)
	if err := <-started; err != nil {
		s.ds.StopSync()
		s.listener.Close()
		return err
	}
	s.TLSAddr = s.tlsListener.Addr().String()

	return nil
}

var ErrClosed = errors.New("strowger: listener has been closed")

func (s *HTTPListener) AddRoute(r *strowger.HTTPRoute) error {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	if s.closed {
		return ErrClosed
	}

	return s.ds.Add(r.Domain, &httpRoute{
		Domain:  r.Domain,
		Service: r.Service,
		TLSCert: r.TLSCert,
		TLSKey:  r.TLSKey,
	})
}

func (s *HTTPListener) RemoveRoute(domain string) error {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	if s.closed {
		return ErrClosed
	}

	return s.ds.Remove(domain)
}

type httpSyncHandler struct {
	l *HTTPListener
}

func (h *httpSyncHandler) Add(data []byte) error {
	r := &httpRoute{}
	if err := json.Unmarshal(data, r); err != nil {
		return err
	}

	if r.TLSCert != "" && r.TLSKey != "" {
		kp, err := tls.X509KeyPair([]byte(r.TLSCert), []byte(r.TLSKey))
		if err != nil {
			return err
		}
		r.keypair = &kp
		r.TLSCert = ""
		r.TLSKey = ""
	}

	h.l.mtx.Lock()
	defer h.l.mtx.Unlock()
	if h.l.closed {
		return ErrClosed
	}
	if _, ok := h.l.domains[r.Domain]; ok {
		return ErrExists
	}

	service := h.l.services[r.Service]
	if service == nil {
		ss, err := h.l.discoverd.NewServiceSet(r.Service)
		if err != nil {
			return err
		}
		service = &httpService{name: r.Service, ss: ss}
		h.l.services[r.Service] = service
	}
	service.refs++
	r.service = service
	h.l.domains[r.Domain] = r

	go h.l.wm.Send(&strowger.Event{Event: "add", ID: r.Domain})
	return nil
}

func (h *httpSyncHandler) Remove(name string) error {
	h.l.mtx.Lock()
	defer h.l.mtx.Unlock()
	r, ok := h.l.domains[name]
	if !ok {
		return ErrNotFound
	}

	r.service.refs--
	if r.service.refs <= 0 {
		r.service.ss.Close()
		delete(h.l.services, r.service.name)
	}

	delete(h.l.domains, name)
	go h.l.wm.Send(&strowger.Event{Event: "remove", ID: name})
	return nil
}

func (s *HTTPListener) serve(started chan<- error) {
	var err error
	s.listener, err = net.Listen("tcp", s.Addr)
	started <- err
	if err != nil {
		return
	}
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			// TODO: log error
			break
		}
		go s.handle(conn, false)
	}
}

func (s *HTTPListener) serveTLS(started chan<- error) {
	var err error
	s.tlsListener, err = net.Listen("tcp", s.TLSAddr)
	started <- err
	if err != nil {
		return
	}
	for {
		conn, err := s.tlsListener.Accept()
		if err != nil {
			// TODO: log error
			break
		}
		go s.handle(conn, true)
	}
}

func (s *HTTPListener) findRouteForHost(host string) *httpRoute {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	// TODO: handle wildcard domains
	backend := s.domains[host]
	return backend
}

func fail(sc *httputil.ServerConn, req *http.Request, code int, msg string) {
	resp := &http.Response{
		StatusCode:    code,
		ProtoMajor:    1,
		ProtoMinor:    0,
		Request:       req,
		Body:          ioutil.NopCloser(bytes.NewBufferString(msg)),
		ContentLength: int64(len(msg)),
	}
	sc.Write(req, resp)
}

func (s *HTTPListener) handle(conn net.Conn, isTLS bool) {
	defer conn.Close()

	var r *httpRoute

	// For TLS, use the SNI hello to determine the domain.
	// At this stage, if we don't find a match, we simply
	// close the connection down.
	if isTLS {
		// Parse out host via SNI first
		vhostConn, err := vhost.TLS(conn)
		if err != nil {
			log.Println("Failed to decode TLS connection", err)
			return
		}
		host := vhostConn.Host()

		// Find a backend for the key
		r = s.findRouteForHost(host)
		if r == nil {
			return
		}
		if r.keypair == nil {
			log.Println("Cannot serve TLS, no certificate defined for this domain")
			return
		}

		// Init a TLS decryptor
		tlscfg := &tls.Config{Certificates: []tls.Certificate{*r.keypair}}
		conn = tls.Server(vhostConn, tlscfg)
	}

	// Decode the first request from the connection
	sc := httputil.NewServerConn(conn, nil)
	req, err := sc.Read()
	if err != nil {
		if err != httputil.ErrPersistEOF {
			// TODO: log error
		}
		return
	}

	// If we do not have a backend yet (unencrypted connection),
	// look at the host header to find one or 404 out.
	if r == nil {
		r = s.findRouteForHost(req.Host)
		if r == nil {
			fail(sc, req, 404, "Not Found")
			return
		}
	}

	r.service.handle(req, sc, isTLS)
}

// A domain served by a listener, associated TLS certs,
// and link to backend service set.
type httpRoute struct {
	Domain  string
	Service string
	TLSCert string
	TLSKey  string

	keypair *tls.Certificate
	service *httpService
}

// A service definition: name, and set of backends.
type httpService struct {
	name string
	ss   discoverd.ServiceSet
	refs int
}

func (s *httpService) getBackend() *httputil.ClientConn {
	for _, addr := range shuffle(s.ss.Addrs()) {
		// TODO: set connection timeout
		backend, err := net.Dial("tcp", addr)
		if err != nil {
			// TODO: log error
			// TODO: limit number of backends tried
			// TODO: temporarily quarantine failing backends
			log.Println("backend error", err)
			continue
		}
		return httputil.NewClientConn(backend, nil)
	}
	// TODO: log no backends found error
	return nil
}

func (s *httpService) handle(req *http.Request, sc *httputil.ServerConn, tls bool) {
	req.Header.Set("X-Request-Start", strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10))
	backend := s.getBackend()
	if backend == nil {
		log.Println("no backend found")
		fail(sc, req, 503, "Service Unavailable")
		return
	}
	defer backend.Close()

	for {
		if req.Method != "GET" && req.Method != "POST" && req.Method != "HEAD" &&
			req.Method != "OPTIONS" && req.Method != "PUT" && req.Method != "DELETE" && req.Method != "TRACE" {
			fail(sc, req, 405, "Method not allowed")
			return
		}

		req.Proto = "HTTP/1.1"
		req.ProtoMajor = 1
		req.ProtoMinor = 1
		delete(req.Header, "Te")
		delete(req.Header, "Transfer-Encoding")

		if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
			// If we aren't the first proxy retain prior
			// X-Forwarded-For information as a comma+space
			// separated list and fold multiple headers into one.
			if prior, ok := req.Header["X-Forwarded-For"]; ok {
				clientIP = strings.Join(prior, ", ") + ", " + clientIP
			}
			req.Header.Set("X-Forwarded-For", clientIP)
		}
		if tls {
			req.Header.Set("X-Forwarded-Proto", "https")
		} else {
			req.Header.Set("X-Forwarded-Proto", "http")
		}
		// TODO: Set X-Forwarded-Port

		if err := backend.Write(req); err != nil {
			log.Println("server write err:", err)
			return
		}
		res, err := backend.Read(req)
		if res != nil {
			if err := sc.Write(req, res); err != nil {
				if err != io.EOF && err != httputil.ErrPersistEOF {
					log.Println("client write err:", err)
					// TODO: log error
				}
				return
			}
		}
		if err != nil {
			if err != io.EOF && err != httputil.ErrPersistEOF {
				log.Println("server read err:", err)
				// TODO: log error
				fail(sc, req, 502, "Bad Gateway")
			}
			return
		}

		// TODO: Proxy HTTP CONNECT? (example: Go RPC over HTTP)
		if res.StatusCode == http.StatusSwitchingProtocols {
			serverW, serverR := backend.Hijack()
			clientW, clientR := sc.Hijack()
			defer serverW.Close()
			done := make(chan struct{})
			go func() {
				serverR.WriteTo(clientW)
				close(done)
			}()
			clientR.WriteTo(serverW)
			<-done
			return
		}

		// TODO: http pipelining
		req, err = sc.Read()
		if err != nil {
			if err != io.EOF && err != httputil.ErrPersistEOF {
				log.Println("client read err:", err)
			}
			return
		}
		req.Header.Set("X-Request-Start", strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10))
	}
}

func shuffle(s []string) []string {
	for i := len(s) - 1; i > 0; i-- {
		j := rand.Intn(i + 1)
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func init() {
	rand.Seed(time.Now().UnixNano())
}
<|endoftext|>


<commit_before>package managesystemagent

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/rancher/fleet/pkg/apis/fleet.cattle.io/v1alpha1"
	rancherv1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	rkev1 "github.com/rancher/rancher/pkg/apis/rke.cattle.io/v1"
	fleetconst "github.com/rancher/rancher/pkg/fleet"
	v3 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	rocontrollers "github.com/rancher/rancher/pkg/generated/controllers/provisioning.cattle.io/v1"
	namespaces "github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/systemtemplate"
	"github.com/rancher/rancher/pkg/wrangler"
	upgradev1 "github.com/rancher/system-upgrade-controller/pkg/apis/upgrade.cattle.io/v1"
	"github.com/rancher/wrangler/pkg/generic"
	"github.com/rancher/wrangler/pkg/gvk"
	"github.com/rancher/wrangler/pkg/name"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

const (
	generationSecretName = "system-agent-upgrade-generation"
)

type handler struct {
	clusterRegistrationTokens v3.ClusterRegistrationTokenCache
}

func Register(ctx context.Context, clients *wrangler.Context) {
	h := &handler{
		clusterRegistrationTokens: clients.Mgmt.ClusterRegistrationToken().Cache(),
	}
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Fleet.Bundle(),
				clients.Provisioning.Cluster(),
				clients.Core.Secret(),
				clients.RBAC.RoleBinding(),
				clients.RBAC.Role()),
		"", "manage-system-agent", h.OnChange, &generic.GeneratingHandlerOptions{
			AllowCrossNamespace: true,
		})
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Mgmt.ManagedChart(),
				clients.Provisioning.Cluster()),
		"", "manage-system-upgrade-controller", h.OnChangeInstallSUC, nil)
}

func (h *handler) OnChange(cluster *rancherv1.Cluster, status rancherv1.ClusterStatus) ([]runtime.Object, rancherv1.ClusterStatus, error) {
	if cluster.Spec.RKEConfig == nil || settings.SystemAgentUpgradeImage.Get() == "" {
		return nil, status, nil
	}

	var (
		secretName = "stv-aggregation"
		result     []runtime.Object
	)

	if cluster.Status.ClusterName == "local" && cluster.Namespace == fleetconst.ClustersLocalNamespace {
		secretName += "-local-"

		token, err := h.clusterRegistrationTokens.Get(cluster.Status.ClusterName, "default-token")
		if err != nil {
			return nil, status, err
		}
		if token.Status.Token == "" {
			return nil, status, fmt.Errorf("token not yet generated for %s/%s", token.Namespace, token.Name)
		}

		digest := sha256.New()
		digest.Write([]byte(settings.InternalServerURL.Get()))
		digest.Write([]byte(token.Status.Token))
		digest.Write([]byte(systemtemplate.InternalCAChecksum()))
		d := digest.Sum(nil)
		secretName += hex.EncodeToString(d[:])[:12]

		result = append(result, &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      secretName,
				Namespace: namespaces.System,
			},
			Data: map[string][]byte{
				"CATTLE_SERVER":      []byte(settings.InternalServerURL.Get()),
				"CATTLE_TOKEN":       []byte(token.Status.Token),
				"CATTLE_CA_CHECKSUM": []byte(systemtemplate.InternalCAChecksum()),
			},
		})
	}

	resources, err := ToResources(installer(cluster.Spec.AgentEnvVars, len(cluster.Spec.RKEConfig.MachineSelectorConfig) == 0, secretName, strconv.Itoa(int(cluster.Spec.RedeploySystemAgentGeneration))))
	if err != nil {
		return nil, status, err
	}

	result = append(result, &v1alpha1.Bundle{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: cluster.Namespace,
			Name:      name.SafeConcatName(cluster.Name, "managed", "system", "agent"),
		},
		Spec: v1alpha1.BundleSpec{
			BundleDeploymentOptions: v1alpha1.BundleDeploymentOptions{
				DefaultNamespace: namespaces.System,
			},
			Resources: resources,
			Targets: []v1alpha1.BundleTarget{
				{
					ClusterName: cluster.Name,
					ClusterSelector: &metav1.LabelSelector{
						MatchExpressions: []metav1.LabelSelectorRequirement{
							{
								Key:      "provisioning.cattle.io/unmanaged-system-agent",
								Operator: metav1.LabelSelectorOpDoesNotExist,
							},
						},
					},
				},
			},
		},
	})

	return result, status, nil
}

func installer(envs []rkev1.EnvVar, allWorkers bool, secretName, generation string) []runtime.Object {
	image := strings.SplitN(settings.SystemAgentUpgradeImage.Get(), ":", 2)
	version := "latest"
	if len(image) == 2 {
		version = image[1]
	}

	var env []corev1.EnvVar
	for _, e := range envs {
		env = append(env, corev1.EnvVar{
			Name:  e.Name,
			Value: e.Value,
		})
	}

	if allWorkers {
		env = append(env, corev1.EnvVar{
			Name:  "CATTLE_ROLE_WORKER",
			Value: "true",
		})
	}

	plan := &upgradev1.Plan{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Plan",
			APIVersion: "upgrade.cattle.io/v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "system-agent-upgrader",
			Namespace: namespaces.System,
			Annotations: map[string]string{
				"upgrade.cattle.io/digest": "spec.upgrade.envs",
			},
		},
		Spec: upgradev1.PlanSpec{
			Concurrency: 10,
			Version:     version,
			Tolerations: []corev1.Toleration{{
				Operator: corev1.TolerationOpExists,
			},
			},
			NodeSelector: &metav1.LabelSelector{
				MatchExpressions: []metav1.LabelSelectorRequirement{
					{
						Key:      corev1.LabelOSStable,
						Operator: metav1.LabelSelectorOpIn,
						Values: []string{
							"linux",
						},
					},
				},
			},
			ServiceAccountName: "system-agent-upgrader",
			Upgrade: &upgradev1.ContainerSpec{
				Image: settings.PrefixPrivateRegistry(image[0]),
				Env:   env,
				EnvFrom: []corev1.EnvFromSource{{
					SecretRef: &corev1.SecretEnvSource{
						LocalObjectReference: corev1.LocalObjectReference{
							Name: secretName,
						},
					},
				}},
			},
		},
	}

	objs := []runtime.Object{
		&corev1.ServiceAccount{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			},
		},
		&rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name: "system-agent-upgrader",
			},
			Rules: []rbacv1.PolicyRule{{
				Verbs:     []string{"get"},
				APIGroups: []string{""},
				Resources: []string{"nodes"},
			}},
		},
		&rbacv1.ClusterRoleBinding{
			ObjectMeta: metav1.ObjectMeta{
				Name: "system-agent-upgrader",
			},
			Subjects: []rbacv1.Subject{{
				Kind:      "ServiceAccount",
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			}},
			RoleRef: rbacv1.RoleRef{
				APIGroup: "rbac.authorization.k8s.io",
				Kind:     "ClusterRole",
				Name:     "system-agent-upgrader",
			},
		},
	}

	if generation != "0" {
		plan.Spec.Secrets = append(plan.Spec.Secrets, upgradev1.SecretSpec{
			Name: generationSecretName,
		})
		objs = append(objs, &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      generationSecretName,
				Namespace: namespaces.System,
			},
			StringData: map[string]string{
				"generation": generation,
			},
		})
	}

	return append([]runtime.Object{plan}, objs...)
}

func ToResources(objs []runtime.Object) (result []v1alpha1.BundleResource, err error) {
	for _, obj := range objs {
		obj = obj.DeepCopyObject()
		if err := gvk.Set(obj); err != nil {
			return nil, fmt.Errorf("failed to set gvk: %w", err)
		}

		typeMeta, err := meta.TypeAccessor(obj)
		if err != nil {
			return nil, err
		}

		meta, err := meta.Accessor(obj)
		if err != nil {
			return nil, err
		}

		data, err := json.Marshal(obj)
		if err != nil {
			return nil, err
		}

		digest := sha256.Sum256(data)
		filename := name.SafeConcatName(typeMeta.GetKind(), meta.GetNamespace(), meta.GetName(), hex.EncodeToString(digest[:])[:12]) + ".yaml"
		result = append(result, v1alpha1.BundleResource{
			Name:    filename,
			Content: string(data),
		})
	}
	return
}
<commit_msg>Add spec.upgrade.envFrom to SUC upgrade watch<commit_after>package managesystemagent

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/rancher/fleet/pkg/apis/fleet.cattle.io/v1alpha1"
	rancherv1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	rkev1 "github.com/rancher/rancher/pkg/apis/rke.cattle.io/v1"
	fleetconst "github.com/rancher/rancher/pkg/fleet"
	v3 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	rocontrollers "github.com/rancher/rancher/pkg/generated/controllers/provisioning.cattle.io/v1"
	namespaces "github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/systemtemplate"
	"github.com/rancher/rancher/pkg/wrangler"
	upgradev1 "github.com/rancher/system-upgrade-controller/pkg/apis/upgrade.cattle.io/v1"
	"github.com/rancher/wrangler/pkg/generic"
	"github.com/rancher/wrangler/pkg/gvk"
	"github.com/rancher/wrangler/pkg/name"
	corev1 "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

const (
	generationSecretName = "system-agent-upgrade-generation"
)

type handler struct {
	clusterRegistrationTokens v3.ClusterRegistrationTokenCache
}

func Register(ctx context.Context, clients *wrangler.Context) {
	h := &handler{
		clusterRegistrationTokens: clients.Mgmt.ClusterRegistrationToken().Cache(),
	}
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Fleet.Bundle(),
				clients.Provisioning.Cluster(),
				clients.Core.Secret(),
				clients.RBAC.RoleBinding(),
				clients.RBAC.Role()),
		"", "manage-system-agent", h.OnChange, &generic.GeneratingHandlerOptions{
			AllowCrossNamespace: true,
		})
	rocontrollers.RegisterClusterGeneratingHandler(ctx, clients.Provisioning.Cluster(),
		clients.Apply.
			WithSetOwnerReference(false, false).
			WithCacheTypes(clients.Mgmt.ManagedChart(),
				clients.Provisioning.Cluster()),
		"", "manage-system-upgrade-controller", h.OnChangeInstallSUC, nil)
}

func (h *handler) OnChange(cluster *rancherv1.Cluster, status rancherv1.ClusterStatus) ([]runtime.Object, rancherv1.ClusterStatus, error) {
	if cluster.Spec.RKEConfig == nil || settings.SystemAgentUpgradeImage.Get() == "" {
		return nil, status, nil
	}

	var (
		secretName = "stv-aggregation"
		result     []runtime.Object
	)

	if cluster.Status.ClusterName == "local" && cluster.Namespace == fleetconst.ClustersLocalNamespace {
		secretName += "-local-"

		token, err := h.clusterRegistrationTokens.Get(cluster.Status.ClusterName, "default-token")
		if err != nil {
			return nil, status, err
		}
		if token.Status.Token == "" {
			return nil, status, fmt.Errorf("token not yet generated for %s/%s", token.Namespace, token.Name)
		}

		digest := sha256.New()
		digest.Write([]byte(settings.InternalServerURL.Get()))
		digest.Write([]byte(token.Status.Token))
		digest.Write([]byte(systemtemplate.InternalCAChecksum()))
		d := digest.Sum(nil)
		secretName += hex.EncodeToString(d[:])[:12]

		result = append(result, &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      secretName,
				Namespace: namespaces.System,
			},
			Data: map[string][]byte{
				"CATTLE_SERVER":      []byte(settings.InternalServerURL.Get()),
				"CATTLE_TOKEN":       []byte(token.Status.Token),
				"CATTLE_CA_CHECKSUM": []byte(systemtemplate.InternalCAChecksum()),
			},
		})
	}

	resources, err := ToResources(installer(cluster.Spec.AgentEnvVars, len(cluster.Spec.RKEConfig.MachineSelectorConfig) == 0, secretName, strconv.Itoa(int(cluster.Spec.RedeploySystemAgentGeneration))))
	if err != nil {
		return nil, status, err
	}

	result = append(result, &v1alpha1.Bundle{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: cluster.Namespace,
			Name:      name.SafeConcatName(cluster.Name, "managed", "system", "agent"),
		},
		Spec: v1alpha1.BundleSpec{
			BundleDeploymentOptions: v1alpha1.BundleDeploymentOptions{
				DefaultNamespace: namespaces.System,
			},
			Resources: resources,
			Targets: []v1alpha1.BundleTarget{
				{
					ClusterName: cluster.Name,
					ClusterSelector: &metav1.LabelSelector{
						MatchExpressions: []metav1.LabelSelectorRequirement{
							{
								Key:      "provisioning.cattle.io/unmanaged-system-agent",
								Operator: metav1.LabelSelectorOpDoesNotExist,
							},
						},
					},
				},
			},
		},
	})

	return result, status, nil
}

func installer(envs []rkev1.EnvVar, allWorkers bool, secretName, generation string) []runtime.Object {
	image := strings.SplitN(settings.SystemAgentUpgradeImage.Get(), ":", 2)
	version := "latest"
	if len(image) == 2 {
		version = image[1]
	}

	var env []corev1.EnvVar
	for _, e := range envs {
		env = append(env, corev1.EnvVar{
			Name:  e.Name,
			Value: e.Value,
		})
	}

	if allWorkers {
		env = append(env, corev1.EnvVar{
			Name:  "CATTLE_ROLE_WORKER",
			Value: "true",
		})
	}

	plan := &upgradev1.Plan{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Plan",
			APIVersion: "upgrade.cattle.io/v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "system-agent-upgrader",
			Namespace: namespaces.System,
			Annotations: map[string]string{
				"upgrade.cattle.io/digest": "spec.upgrade.envs,spec.upgrade.envFrom",
			},
		},
		Spec: upgradev1.PlanSpec{
			Concurrency: 10,
			Version:     version,
			Tolerations: []corev1.Toleration{{
				Operator: corev1.TolerationOpExists,
			},
			},
			NodeSelector: &metav1.LabelSelector{
				MatchExpressions: []metav1.LabelSelectorRequirement{
					{
						Key:      corev1.LabelOSStable,
						Operator: metav1.LabelSelectorOpIn,
						Values: []string{
							"linux",
						},
					},
				},
			},
			ServiceAccountName: "system-agent-upgrader",
			Upgrade: &upgradev1.ContainerSpec{
				Image: settings.PrefixPrivateRegistry(image[0]),
				Env:   env,
				EnvFrom: []corev1.EnvFromSource{{
					SecretRef: &corev1.SecretEnvSource{
						LocalObjectReference: corev1.LocalObjectReference{
							Name: secretName,
						},
					},
				}},
			},
		},
	}

	objs := []runtime.Object{
		&corev1.ServiceAccount{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			},
		},
		&rbacv1.ClusterRole{
			ObjectMeta: metav1.ObjectMeta{
				Name: "system-agent-upgrader",
			},
			Rules: []rbacv1.PolicyRule{{
				Verbs:     []string{"get"},
				APIGroups: []string{""},
				Resources: []string{"nodes"},
			}},
		},
		&rbacv1.ClusterRoleBinding{
			ObjectMeta: metav1.ObjectMeta{
				Name: "system-agent-upgrader",
			},
			Subjects: []rbacv1.Subject{{
				Kind:      "ServiceAccount",
				Name:      "system-agent-upgrader",
				Namespace: namespaces.System,
			}},
			RoleRef: rbacv1.RoleRef{
				APIGroup: "rbac.authorization.k8s.io",
				Kind:     "ClusterRole",
				Name:     "system-agent-upgrader",
			},
		},
	}

	if generation != "0" {
		plan.Spec.Secrets = append(plan.Spec.Secrets, upgradev1.SecretSpec{
			Name: generationSecretName,
		})
		objs = append(objs, &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      generationSecretName,
				Namespace: namespaces.System,
			},
			StringData: map[string]string{
				"generation": generation,
			},
		})
	}

	return append([]runtime.Object{plan}, objs...)
}

func ToResources(objs []runtime.Object) (result []v1alpha1.BundleResource, err error) {
	for _, obj := range objs {
		obj = obj.DeepCopyObject()
		if err := gvk.Set(obj); err != nil {
			return nil, fmt.Errorf("failed to set gvk: %w", err)
		}

		typeMeta, err := meta.TypeAccessor(obj)
		if err != nil {
			return nil, err
		}

		meta, err := meta.Accessor(obj)
		if err != nil {
			return nil, err
		}

		data, err := json.Marshal(obj)
		if err != nil {
			return nil, err
		}

		digest := sha256.Sum256(data)
		filename := name.SafeConcatName(typeMeta.GetKind(), meta.GetNamespace(), meta.GetName(), hex.EncodeToString(digest[:])[:12]) + ".yaml"
		result = append(result, v1alpha1.BundleResource{
			Name:    filename,
			Content: string(data),
		})
	}
	return
}
<|endoftext|>