Implemented flushing interval (#23)

* Implemented flushing interval

When proxying streaming responses, it would not flush the response writer buffer until some seemingly random point (maybe the number of bytes?). This makes it flush every 1 second by default, but with a configurable interval.

* flushing CHANGELOG

* gofmt and goimports

CHANGELOG.md

@@ -10,6 +10,8 @@
 - [#21](https://github.com/pusher/oauth2_proxy/pull/21) Docker Improvement (@yaegashi)
   - Move Docker base image from debian to alpine
   - Install ca-certificates in docker image
+- [#23](https://github.com/pusher/oauth2_proxy/pull/21) Flushed streaming responses
+  - Long-running upstream responses will get flushed every <timeperiod> (1 second by default)
 - [#24](https://github.com/pusher/oauth2_proxy/pull/24) Redirect fix (@agentgonzo)
   - After a successful login, you will be redirected to your original URL rather than /
 

logging_handler.go

@@ -76,6 +76,12 @@ func (l *responseLogger) Size() int {
 	return l.size
 }
 
+func (l *responseLogger) Flush() {
+	if flusher, ok := l.w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
 // logMessageData is the container for all values that are available as variables in the request logging format.
 // All values are pre-formatted strings so it is easy to use them in the format string.
 type logMessageData struct {

main.go

@@ -43,6 +43,7 @@ func main() {
 	flagSet.Bool("skip-provider-button", false, "will skip sign-in-page to directly reach the next step: oauth/start")
 	flagSet.Bool("skip-auth-preflight", false, "will skip authentication for OPTIONS requests")
 	flagSet.Bool("ssl-insecure-skip-verify", false, "skip validation of certificates presented when using HTTPS")
+	flagSet.Duration("flush-interval", time.Duration(1)*time.Second, "period between response flushing when streaming responses")
 
 	flagSet.Var(&emailDomains, "email-domain", "authenticate emails with the specified domain (may be given multiple times). Use * to authenticate any email")
 	flagSet.String("azure-tenant", "common", "go to a tenant-specific or common (tenant-independent) endpoint.")

oauthproxy.go

@@ -110,8 +110,10 @@ func (u *UpstreamProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 // NewReverseProxy creates a new reverse proxy for proxying requests to upstream
 // servers
-func NewReverseProxy(target *url.URL) (proxy *httputil.ReverseProxy) {
-	return httputil.NewSingleHostReverseProxy(target)
+func NewReverseProxy(target *url.URL, flushInterval time.Duration) (proxy *httputil.ReverseProxy) {
+	proxy = httputil.NewSingleHostReverseProxy(target)
+	proxy.FlushInterval = flushInterval
+	return proxy
 }
 
 func setProxyUpstreamHostHeader(proxy *httputil.ReverseProxy, target *url.URL) {
@@ -154,7 +156,7 @@ func NewOAuthProxy(opts *Options, validator func(string) bool) *OAuthProxy {
 		case httpScheme, httpsScheme:
 			u.Path = ""
 			log.Printf("mapping path %q => upstream %q", path, u)
-			proxy := NewReverseProxy(u)
+			proxy := NewReverseProxy(u, opts.FlushInterval)
 			if !opts.PassHostHeader {
 				setProxyUpstreamHostHeader(proxy, u)
 			} else {

oauthproxy_test.go

@@ -39,7 +39,7 @@ func TestNewReverseProxy(t *testing.T) {
 	backendHost := net.JoinHostPort(backendHostname, backendPort)
 	proxyURL, _ := url.Parse(backendURL.Scheme + "://" + backendHost + "/")
 
-	proxyHandler := NewReverseProxy(proxyURL)
+	proxyHandler := NewReverseProxy(proxyURL, time.Second)
 	setProxyUpstreamHostHeader(proxyHandler, proxyURL)
 	frontend := httptest.NewServer(proxyHandler)
 	defer frontend.Close()
@@ -61,7 +61,7 @@ func TestEncodedSlashes(t *testing.T) {
 	defer backend.Close()
 
 	b, _ := url.Parse(backend.URL)
-	proxyHandler := NewReverseProxy(b)
+	proxyHandler := NewReverseProxy(b, time.Second)
 	setProxyDirector(proxyHandler)
 	frontend := httptest.NewServer(proxyHandler)
 	defer frontend.Close()

options.go

@@ -64,6 +64,7 @@ type Options struct {
 	SetAuthorization      bool          `flag:"set-authorization-header" cfg:"set_authorization_header"`
 	PassAuthorization     bool          `flag:"pass-authorization-header" cfg:"pass_authorization_header"`
 	SkipAuthPreflight     bool          `flag:"skip-auth-preflight" cfg:"skip_auth_preflight"`
+	FlushInterval         time.Duration `flag:"flush-interval" cfg:"flush_interval"`
 
 	// These options allow for other providers besides Google, with
 	// potential overrides.