Merge pull request #81 from netlify/safedial

Add SafeDialContext method

Author: biilmann

graceful/graceful.go

@@ -0,0 +1,107 @@
+package graceful
+
+import (
+	"context"
+	"errors"
+	"net"
+	"net/http"
+	"os"
+	"os/signal"
+	"sync/atomic"
+	"syscall"
+	"time"
+
+	"github.com/sirupsen/logrus"
+)
+
+var DefaultShutdownTimeout = time.Second * 60
+
+const shutdown uint32 = 1
+
+type GracefulServer struct {
+	server   *http.Server
+	listener net.Listener
+	log      *logrus.Entry
+
+	exit chan struct{}
+
+	URL             string
+	state           uint32
+	ShutdownTimeout time.Duration
+	ShutdownError   error
+}
+
+func NewGracefulServer(handler http.Handler, log *logrus.Entry) *GracefulServer {
+	log.Warn("NewGracefulServer is deprecated, see https://github.com/netlify/netlify-commons/pull/72")
+	return &GracefulServer{
+		server:          &http.Server{Handler: handler},
+		log:             log,
+		listener:        nil,
+		exit:            make(chan struct{}),
+		ShutdownTimeout: DefaultShutdownTimeout,
+	}
+}
+
+func (svr *GracefulServer) Bind(addr string) error {
+	l, err := net.Listen("tcp", addr)
+	if err != nil {
+		return err
+	}
+	svr.URL = "http://" + l.Addr().String()
+	svr.listener = l
+	return nil
+}
+
+func (svr *GracefulServer) Listen() error {
+	go svr.listenForShutdownSignal()
+	serveErr := svr.server.Serve(svr.listener)
+	if serveErr != http.ErrServerClosed {
+		svr.log.WithError(serveErr).Warn("Error while running server")
+		return serveErr
+	}
+
+	<-svr.exit
+
+	return svr.ShutdownError
+}
+
+func (svr *GracefulServer) listenForShutdownSignal() {
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
+	sig := <-c
+	svr.log.Infof("Triggering shutdown from signal %s", sig)
+	svr.Shutdown()
+}
+
+func (svr *GracefulServer) ListenAndServe(addr string) error {
+	if svr.listener != nil {
+		return errors.New("The listener has already started, don't call Bind first")
+	}
+	if err := svr.Bind(addr); err != nil {
+		return err
+	}
+
+	return svr.Listen()
+}
+
+func (svr *GracefulServer) Shutdown() error {
+	if atomic.SwapUint32(&svr.state, shutdown) == shutdown {
+		svr.log.Debug("Calling shutdown on already shutdown server, ignoring")
+		return nil
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), svr.ShutdownTimeout)
+	defer cancel()
+
+	svr.log.Infof("Triggering shutdown, in at most %s ", svr.ShutdownTimeout.String())
+	shutErr := svr.server.Shutdown(ctx)
+	if shutErr == context.DeadlineExceeded {
+		svr.log.WithError(shutErr).Warnf("Forcing a shutdown after waiting %s", svr.ShutdownTimeout.String())
+		shutErr = svr.server.Close()
+	}
+
+	svr.ShutdownError = shutErr
+	close(svr.exit)
+
+	return shutErr
+}

graceful/graceful_test.go

@@ -0,0 +1,94 @@
+package graceful
+
+import (
+	"fmt"
+	"net/http"
+	"net/url"
+	"reflect"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestStartAndStop(t *testing.T) {
+	orTimeout := func(c chan bool, i int, msg string) {
+		select {
+		case <-c:
+		case <-time.After(time.Duration(i) * time.Second):
+			assert.FailNow(t, msg)
+		}
+	}
+
+	gotRequest := make(chan bool)
+	clearRequest := make(chan bool)
+	stoppedServer := make(chan bool)
+	finishedListening := make(chan bool)
+
+	var finished bool
+
+	oh := func(w http.ResponseWriter, r *http.Request) {
+		// trigger that we got the request
+		close(gotRequest)
+		// wait for a clear on that request
+		orTimeout(clearRequest, 2, "waiting for request to be cleared")
+		finished = true
+	}
+
+	svr := NewGracefulServer(http.HandlerFunc(oh), logrus.WithField("testing", true))
+	require.NoError(t, svr.Bind("127.0.0.1:0"))
+
+	go func() {
+		assert.NoError(t, svr.Listen())
+		close(finishedListening)
+	}()
+
+	// make a request
+	go func() {
+		rsp, err := http.Get(svr.URL + "/something")
+		require.NoError(t, err)
+		assert.Equal(t, http.StatusOK, rsp.StatusCode)
+	}()
+
+	// wait for the origin to get the request
+	orTimeout(gotRequest, 1, "didn't get the original request in time")
+
+	// initate a shutdown
+	go func() {
+		assert.NoError(t, svr.Shutdown())
+		close(stoppedServer)
+	}()
+
+	<-time.After(time.Second)
+
+	// make a second request ~ should be bounced
+	rsp, err := http.Get(svr.URL + "/something")
+	switch e := err.(type) {
+	case *url.Error:
+		assert.True(t, strings.Contains(e.Error(), "connection refused"))
+	default:
+		assert.Fail(t, fmt.Sprintf("unknown type: %v", reflect.TypeOf(err)))
+	}
+	assert.Nil(t, rsp)
+
+	// finish the first request
+	close(clearRequest)
+
+	// wait for server to close
+	orTimeout(stoppedServer, 1, "didn't stop server in time")
+
+	assert.True(t, finished)
+	orTimeout(finishedListening, 1, "didn't actually stop the server in time")
+}
+
+func TestDoubleShutdown(t *testing.T) {
+	logrus.SetLevel(logrus.DebugLevel)
+	oh := func(w http.ResponseWriter, r *http.Request) {}
+	svr := NewGracefulServer(http.HandlerFunc(oh), logrus.WithField("testing", true))
+
+	assert.NoError(t, svr.Shutdown())
+	assert.NoError(t, svr.Shutdown())
+}

http/http.go

@@ -0,0 +1,52 @@
+package http
+
+import (
+	"context"
+	"errors"
+	"net"
+)
+
+var privateIPBlocks []*net.IPNet
+
+func init() {
+	for _, cidr := range []string{
+		"127.0.0.0/8",    // IPv4 loopback
+		"10.0.0.0/8",     // RFC1918
+		"100.64.0.0/10",  // RFC6598
+		"172.16.0.0/12",  // RFC1918
+		"192.0.0.0/24",   // RFC6890
+		"192.168.0.0/16", // RFC1918
+		"169.254.0.0/16", // RFC3927
+		"::1/128",        // IPv6 loopback
+		"fe80::/10",      // IPv6 link-local
+		"fc00::/7",       // IPv6 unique local addr
+	} {
+		_, block, _ := net.ParseCIDR(cidr)
+		privateIPBlocks = append(privateIPBlocks, block)
+	}
+}
+
+func isPrivateIP(ip net.IP) bool {
+	for _, block := range privateIPBlocks {
+		if block.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+func isLocalAddress(addr string) bool {
+	ip := net.ParseIP(addr)
+	return isPrivateIP(ip)
+}
+
+// SafeDialContext exchanges a DialContext for a SafeDialContext that will never dial a reserved IP range
+func SafeDialContext(dialContext func(ctx context.Context, network, addr string) (net.Conn, error)) func(ctx context.Context, network, addr string) (net.Conn, error) {
+	return func(ctx context.Context, network, addr string) (net.Conn, error) {
+		if isLocalAddress(addr) {
+			return nil, errors.New("Connection to local network address denied")
+		}
+
+		return dialContext(ctx, network, addr)
+	}
+}

http/http_test.go

@@ -0,0 +1,16 @@
+package http
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsLocalAddress(t *testing.T) {
+	assert.False(t, isLocalAddress("216.58.194.206"))
+	assert.True(t, isLocalAddress("127.0.0.1"))
+	assert.True(t, isLocalAddress("10.0.0.1"))
+	assert.True(t, isLocalAddress("192.168.0.1"))
+	assert.True(t, isLocalAddress("172.16.0.0"))
+	assert.True(t, isLocalAddress("169.254.169.254"))
+}