initial commit

tls.go

@@ -0,0 +1,94 @@
+package graceful
+
+import (
+	"crypto/tls"
+	"errors"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+)
+
+// TLSListener is a TLS wrapper for net.TCPListener
+type TLSListener struct {
+	// protects conns
+	sync.Mutex
+
+	*net.TCPListener
+
+	config    *tls.Config
+	conns     map[net.Conn]struct{}
+	keepAlive time.Duration
+}
+
+// TLSConn is a net.Conn wrapper
+type TLSConn struct {
+	net.Conn
+	*TLSListener
+}
+
+// Close wraps the net.Conn.Close method and removes the conn from the listener map
+func (c *TLSConn) Close() error {
+	c.RemoveConn(c.Conn)
+	return c.Conn.Close()
+}
+
+// NewTLS wraps TLSListener around net.TCPListener
+func NewTLS(laddr string, config *tls.Config, keepAlive time.Duration) (*TLSListener, error) {
+	if config == nil || len(config.Certificates) == 0 {
+		return nil, errors.New("NewTLS: no certificates in configuration")
+	}
+	l, err := net.Listen(network, laddr)
+	if err != nil {
+		return nil, err
+	}
+	ltcp, ok := l.(*net.TCPListener)
+	if !ok {
+		return nil, ErrInterfaceConversion
+	}
+	return &TLSListener{
+		TCPListener: ltcp,
+		config:      config,
+		conns:       make(map[net.Conn]struct{}),
+		keepAlive:   keepAlive,
+	}, nil
+}
+
+// Accept wraps the net.TCPListener.AcceptTCP method
+func (l *TLSListener) Accept() (net.Conn, error) {
+	conn, err := l.AcceptTCP()
+	if err != nil {
+		return conn, err
+	}
+	if l.keepAlive > 0 {
+		conn.SetKeepAlive(true)
+		conn.SetKeepAlivePeriod(l.keepAlive)
+	}
+	tlsconn := &TLSConn{tls.Server(conn, l.config), l}
+	l.Lock()
+	defer l.Unlock()
+	l.conns[tlsconn] = struct{}{}
+	return tlsconn, nil
+}
+
+// TLSConnState implements HTTP connection hijacking
+func (l *TLSListener) TLSConnState(c net.Conn, cs http.ConnState) {
+	if cs == http.StateClosed || cs == http.StateHijacked {
+		l.RemoveConn(c)
+	}
+}
+
+// RemoveConn removes a connection from the map
+func (l *TLSListener) RemoveConn(c net.Conn) {
+	l.Lock()
+	defer l.Unlock()
+	delete(l.conns, c)
+}
+
+// Stop stops the listener and closes all open connections
+func (l *TLSListener) Stop() {
+	l.Close()
+	for c := range l.conns {
+		c.Close()
+	}
+}