upstream.go 4.82 KB
Newer Older
1 2 3 4 5 6
package proxy

import (
	"io"
	"io/ioutil"
	"net/http"
7 8 9
	"net/url"
	"strconv"
	"strings"
10
	"time"
11 12

	"github.com/mholt/caddy/config/parse"
13 14
)

15 16
var (
	supportedPolicies map[string]func() Policy = make(map[string]func() Policy)
Austin's avatar
Austin committed
17
	proxyHeaders      http.Header              = make(http.Header)
18
)
Austin's avatar
Austin committed
19

20 21
type staticUpstream struct {
	from   string
22 23 24 25 26 27 28 29 30
	Hosts  HostPool
	Policy Policy

	FailTimeout time.Duration
	MaxFails    int32
	HealthCheck struct {
		Path     string
		Interval time.Duration
	}
31
	WithoutPathPrefix string
32 33
}

34
// NewStaticUpstreams parses the configuration input and sets up
35 36 37 38 39 40 41 42 43 44 45
// static upstreams for the proxy middleware.
func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) {
	var upstreams []Upstream
	for c.Next() {
		upstream := &staticUpstream{
			from:        "",
			Hosts:       nil,
			Policy:      &Random{},
			FailTimeout: 10 * time.Second,
			MaxFails:    1,
		}
46

47 48 49 50 51 52 53 54 55 56 57 58 59 60
		if !c.Args(&upstream.from) {
			return upstreams, c.ArgErr()
		}
		to := c.RemainingArgs()
		if len(to) == 0 {
			return upstreams, c.ArgErr()
		}

		for c.NextBlock() {
			switch c.Val() {
			case "policy":
				if !c.NextArg() {
					return upstreams, c.ArgErr()
				}
Austin's avatar
Austin committed
61

Austin's avatar
Austin committed
62 63 64
				if policyCreateFunc, ok := supportedPolicies[c.Val()]; ok {
					upstream.Policy = policyCreateFunc()
				} else {
65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
					return upstreams, c.ArgErr()
				}
			case "fail_timeout":
				if !c.NextArg() {
					return upstreams, c.ArgErr()
				}
				if dur, err := time.ParseDuration(c.Val()); err == nil {
					upstream.FailTimeout = dur
				} else {
					return upstreams, err
				}
			case "max_fails":
				if !c.NextArg() {
					return upstreams, c.ArgErr()
				}
				if n, err := strconv.Atoi(c.Val()); err == nil {
					upstream.MaxFails = int32(n)
				} else {
					return upstreams, err
				}
			case "health_check":
				if !c.NextArg() {
					return upstreams, c.ArgErr()
				}
				upstream.HealthCheck.Path = c.Val()
				upstream.HealthCheck.Interval = 30 * time.Second
				if c.NextArg() {
					if dur, err := time.ParseDuration(c.Val()); err == nil {
						upstream.HealthCheck.Interval = dur
					} else {
						return upstreams, err
					}
				}
			case "proxy_header":
				var header, value string
				if !c.Args(&header, &value) {
					return upstreams, c.ArgErr()
				}
Austin's avatar
Austin committed
103
				proxyHeaders.Add(header, value)
104
			case "websocket":
Austin's avatar
Austin committed
105 106
				proxyHeaders.Add("Connection", "{>Connection}")
				proxyHeaders.Add("Upgrade", "{>Upgrade}")
107 108 109 110
			case "without":
				if !c.NextArg() {
					return upstreams, c.ArgErr()
				}
111
				upstream.WithoutPathPrefix = c.Val()
112 113
			default:
				return upstreams, c.Errf("unknown property '%s'", c.Val())
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
			}
		}

		upstream.Hosts = make([]*UpstreamHost, len(to))
		for i, host := range to {
			if !strings.HasPrefix(host, "http") {
				host = "http://" + host
			}
			uh := &UpstreamHost{
				Name:         host,
				Conns:        0,
				Fails:        0,
				FailTimeout:  upstream.FailTimeout,
				Unhealthy:    false,
				ExtraHeaders: proxyHeaders,
				CheckDown: func(upstream *staticUpstream) UpstreamHostDownFunc {
					return func(uh *UpstreamHost) bool {
						if uh.Unhealthy {
							return true
						}
						if uh.Fails >= upstream.MaxFails &&
							upstream.MaxFails != 0 {
							return true
						}
						return false
					}
				}(upstream),
141
				WithoutPathPrefix: upstream.WithoutPathPrefix,
142
			}
143
			if baseURL, err := url.Parse(uh.Name); err == nil {
144
				uh.ReverseProxy = NewSingleHostReverseProxy(baseURL, uh.WithoutPathPrefix)
145 146 147 148 149 150 151 152 153 154 155 156 157 158
			} else {
				return upstreams, err
			}
			upstream.Hosts[i] = uh
		}

		if upstream.HealthCheck.Path != "" {
			go upstream.HealthCheckWorker(nil)
		}
		upstreams = append(upstreams, upstream)
	}
	return upstreams, nil
}

Austin's avatar
Austin committed
159
// RegisterPolicy adds a custom policy to the proxy.
Austin's avatar
Austin committed
160 161
func RegisterPolicy(name string, policy func() Policy) {
	supportedPolicies[name] = policy
Austin's avatar
Austin committed
162 163
}

164 165
func (u *staticUpstream) From() string {
	return u.from
166 167
}

168
func (u *staticUpstream) healthCheck() {
169
	for _, host := range u.Hosts {
170 171
		hostURL := host.Name + u.HealthCheck.Path
		if r, err := http.Get(hostURL); err == nil {
172 173 174 175 176 177 178 179 180
			io.Copy(ioutil.Discard, r.Body)
			r.Body.Close()
			host.Unhealthy = r.StatusCode < 200 || r.StatusCode >= 400
		} else {
			host.Unhealthy = true
		}
	}
}

181
func (u *staticUpstream) HealthCheckWorker(stop chan struct{}) {
182 183 184 185 186 187 188 189 190 191 192 193 194 195
	ticker := time.NewTicker(u.HealthCheck.Interval)
	u.healthCheck()
	for {
		select {
		case <-ticker.C:
			u.healthCheck()
		case <-stop:
			// TODO: the library should provide a stop channel and global
			// waitgroup to allow goroutines started by plugins a chance
			// to clean themselves up.
		}
	}
}

196
func (u *staticUpstream) Select() *UpstreamHost {
197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217
	pool := u.Hosts
	if len(pool) == 1 {
		if pool[0].Down() {
			return nil
		}
		return pool[0]
	}
	allDown := true
	for _, host := range pool {
		if !host.Down() {
			allDown = false
			break
		}
	}
	if allDown {
		return nil
	}

	if u.Policy == nil {
		return (&Random{}).Select(pool)
	}
218
	return u.Policy.Select(pool)
219
}