workaround for armon/go-proxyproto#1

proxyproto_workaround.go

package main

import (
	"errors"
	"net"
	"net/http"
	"sync"
)

// This file is a silly workaround for an issue caused by an interaction with
// how the proxyproto library works and how http.Serve works.  The proxyproto
// library causes `conn.RemoteAddr` to be a blocking operation. `http.Serve`
// calls this method (in `srv.newConn`) between `Accept` and spawning a new
// goroutine. So this can cause the entire server to block!  So we wrap the
// listener in such a way that we are able to start a new accepting goroutine
// whenever one call to accept finishes.

type singleUseListener struct {
	l net.Listener
	// XXX is mutex needed? only one goroutine gets each singleUseListener.
	// technically srv.Serve could do something fancy but we're assuming
	// a lot about its implementation already.
	mu           sync.Mutex
	done         bool
	acceptErrors chan error
}

func isTemporaryError(e error) bool {
	if e == nil {
		return false
	}
	ne, ok := e.(net.Error)
	if !ok {
		return false
	}
	return ne.Temporary()
}

func (s *singleUseListener) Accept() (net.Conn, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.done {
		return nil, errors.New("single use is done")
	}
	conn, acceptErr := s.l.Accept()

	// srv.Serve will interpret certain "temporary" errors as meaning it should
	// sleep and retry again.  Allow this Listener to be used again if it was
	// one of those errors.
	if isTemporaryError(acceptErr) {
		return conn, acceptErr
	}

	s.done = true
	s.acceptErrors <- acceptErr
	return conn, acceptErr
}

func (s *singleUseListener) Close() error {
	// Do nothing: we *don't* want to close the underlying listener when
	// http.Serve returns!
	return nil
}

func (s *singleUseListener) Addr() net.Addr {
	return s.l.Addr()
}

func serveHTTP(listener net.Listener, handler http.Handler) error {
	srv := &http.Server{Handler: handler}

	// This channel serves two purposes: getting real Accept errors so we can
	// return them from serveHTTP, and letting us know when Accept has finished so
	// that we can start the next serving goroutine.
	acceptErrors := make(chan error, 1)
	// Prime the channel so that we start a serving goroutine the first time we
	// hit the loop.
	acceptErrors <- nil

	// srv.Serve ordinarily returns its listener when it closes, but
	// singleUseListener stops that (by making Close into a no-op); so let's do it
	// ourselves (just once).  Ignoring errors matches srv.Serve.
	defer func() { _ = listener.Close() }()

	for {
		// Wait for a call to the wrapped Accept() to finish.
		acceptErr := <-acceptErrors
		if acceptErr != nil {
			// There's an error; return it.  (The last goroutine that we spawned will also
			// return this error and conclude.)
			return acceptErr
		}

		// No error! Let's start a new serving goroutine that can run in parallel with the
		// rest of the last connection's srv.Serve.
		go func() {
			// Why ignore errors?
			//
			// We assume, based on reading the 1.5.1 source, that srv.Serve only ever
			// returns errors that it gets from l.Accept.
			//
			// In this scenario, these errors will be of two types: either an error
			// returned by the wrapped listener's Accept, in which case we get it
			// through the channel above and return it; or the "single use is done"
			// error, in which case we should not report it to callers. In either
			// case, ignoring the error from http.Serve is correct.
			_ = srv.Serve(&singleUseListener{l: listener, acceptErrors: acceptErrors})
		}()
	}
}